diff --git a/components/spider-scheduler/src/config.rs b/components/spider-scheduler/src/config.rs index db585d3ae..0aa973e91 100644 --- a/components/spider-scheduler/src/config.rs +++ b/components/spider-scheduler/src/config.rs @@ -8,7 +8,6 @@ use spider_utils::config::EndpointConfig; use crate::core::SchedulerCore; use crate::core_impl::RoundRobinConfig; -use crate::dispatch_queue::DispatchQueueSink; use crate::runtime::RuntimeConfig; use crate::storage_client::SchedulerStorageClient; @@ -52,34 +51,16 @@ impl SchedulerConfig { /// # Type Parameters /// /// * `SchedulerStorageClientType` - The storage client the core polls and registers through. - /// * `DispatchQueueSinkType` - The dispatch sink that task assignments are written to. /// /// # Returns /// /// A boxed [`SchedulerCore`] configured by the selected variant. #[must_use] - pub fn make_core< - SchedulerStorageClientType: SchedulerStorageClient + 'static, - DispatchQueueSinkType: DispatchQueueSink + 'static, - >( + pub fn make_core( self, - ) -> Box< - dyn SchedulerCore, - > { + ) -> Box> { match self { - Self::RoundRobin(config) => { - Box::new(config.make_core::()) - } - } - } - - /// # Returns - /// - /// The dispatch queue capacity of the selected variant. - #[must_use] - pub const fn dispatch_queue_capacity(&self) -> std::num::NonZeroUsize { - match self { - Self::RoundRobin(config) => config.dispatch_queue_capacity, + Self::RoundRobin(config) => Box::new(config.make_core::()), } } } diff --git a/components/spider-scheduler/src/core.rs b/components/spider-scheduler/src/core.rs index 4a2bda858..dc2ae1b26 100644 --- a/components/spider-scheduler/src/core.rs +++ b/components/spider-scheduler/src/core.rs @@ -6,7 +6,7 @@ use std::sync::atomic::AtomicU64; use async_trait::async_trait; use spider_core::types::id::TaskAssignmentId; -use crate::dispatch_queue::DispatchQueueSink; +use crate::dispatch_queue::SharedDispatchQueueHandle; use crate::error::SchedulerError; use crate::storage_client::SchedulerStorageClient; use crate::types::TaskAssignment; @@ -39,22 +39,26 @@ impl TaskAssignmentIdIssuer { /// An abstracted core for a scheduling algorithm. /// /// A core owns its decision loop: it polls the inbound queue through a [`SchedulerStorageClient`], -/// applies its algorithm (reading storage as needed for placement), and writes assignments to a -/// [`DispatchQueueSink`]. Modeling the algorithm as a trait lets different scheduling strategies -/// share the same runtime entry point. +/// applies its algorithm (reading storage as needed for placement), and writes assignments into the +/// dispatch structure it owns and constructs. Modeling the algorithm as a trait lets different +/// scheduling strategies share the same runtime entry point while each picks the dispatch structure +/// its algorithm needs. #[async_trait] pub trait SchedulerCore: Send { - /// The dispatch sink the core writes assignments to. - type Sink: DispatchQueueSink; - /// The storage client used by the core to poll and read for placement decisions. type StorageClient: SchedulerStorageClient; + /// # Returns + /// + /// A handle over the core's dispatch structure, for the execution-manager-facing service to + /// drain. + fn get_dispatch_queue_handle(&self) -> SharedDispatchQueueHandle; + /// Runs the scheduling loop until `cancellation_token` is triggered. /// /// The core polls the inbound queue through `storage_client`, applies its scheduling algorithm, - /// and writes assignments to `sink`, repeating until `cancellation_token` is fired, at which - /// point it returns. + /// and writes assignments into its dispatch structure, repeating until `cancellation_token` is + /// fired, at which point it returns. /// /// The implementation does not need to fire the cancellation token when exiting on error: this /// cancellation is handled by the runtime. @@ -63,7 +67,6 @@ pub trait SchedulerCore: Send { /// /// * `storage_client` - The storage client used to poll the inbound queue and read state for /// placement. - /// * `sink` - The dispatch sink that assignments are written to. /// * `reschedule_queue_reader` - The reader side of the re-schedule queue, delivering task /// assignments returned for re-placement when an execution manager is lost. /// * `id_issuer` - The single-source ID issuer for creating globally unique IDs for task @@ -76,7 +79,6 @@ pub trait SchedulerCore: Send { async fn run( self: Box, storage_client: Self::StorageClient, - sink: Self::Sink, reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver, id_issuer: TaskAssignmentIdIssuer, cancellation_token: tokio_util::sync::CancellationToken, diff --git a/components/spider-scheduler/src/core_impl/round_robin/implementation.rs b/components/spider-scheduler/src/core_impl/round_robin/implementation.rs index dd9a573a7..6d5a362ad 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/implementation.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/implementation.rs @@ -6,6 +6,7 @@ use std::collections::HashSet; use std::collections::VecDeque; use std::num::NonZeroU64; use std::num::NonZeroUsize; +use std::sync::Arc; use std::time::Duration; use std::time::Instant; @@ -18,7 +19,6 @@ use spider_core::types::id::TaskId; use tokio::select; use tokio_util::sync::CancellationToken; -use crate::DispatchQueueSink; use crate::InboundEntry; use crate::SchedulerCore; use crate::SchedulerError; @@ -26,6 +26,9 @@ use crate::SchedulerStorageClient; use crate::StorageClientError; use crate::TaskAssignment; use crate::core::TaskAssignmentIdIssuer; +use crate::dispatch_queue::DispatchQueueWriter; +use crate::dispatch_queue::SharedDispatchQueueHandle; +use crate::dispatch_queue::create_dispatch_queue; /// The configuration of the round-robin scheduler core. #[derive(Clone, Debug, Deserialize)] @@ -64,20 +67,20 @@ impl RoundRobinConfig { /// # Type Parameters /// /// * `SchedulerStorageClientType` - The storage client used to poll the inbound queue. - /// * `DispatchQueueSinkType` - The dispatch sink that task assignments are written to. /// /// # Returns /// - /// A newly created round-robin scheduler core. + /// A newly created round-robin scheduler core, owning a freshly created dispatch queue. #[must_use] - pub const fn make_core< - SchedulerStorageClientType: SchedulerStorageClient + 'static, - DispatchQueueSinkType: DispatchQueueSink, - >( + pub fn make_core( self, - ) -> RoundRobinCore { + ) -> RoundRobinCore { + let (dispatch_queue_writer, dispatch_queue_reader) = + create_dispatch_queue(self.dispatch_queue_capacity.get(), SessionId::default()); RoundRobinCore { config: self, + dispatch_queue_writer, + dispatch_queue_reader: Arc::new(dispatch_queue_reader), _marker: std::marker::PhantomData, } } @@ -92,40 +95,43 @@ impl RoundRobinConfig { /// # Type Parameters /// /// * `SchedulerStorageClientType` - The storage client used to poll the inbound queue. -/// * `DispatchQueueSinkType` - The dispatch sink that task assignments are written to. -pub struct RoundRobinCore< - SchedulerStorageClientType: SchedulerStorageClient + 'static, - DispatchQueueSinkType: DispatchQueueSink, -> { +pub struct RoundRobinCore { config: RoundRobinConfig, - _marker: std::marker::PhantomData<(SchedulerStorageClientType, DispatchQueueSinkType)>, + dispatch_queue_writer: DispatchQueueWriter, + dispatch_queue_reader: SharedDispatchQueueHandle, + _marker: std::marker::PhantomData, } #[async_trait] -impl< - SchedulerStorageClientType: SchedulerStorageClient + 'static, - DispatchQueueSinkType: DispatchQueueSink, -> SchedulerCore for RoundRobinCore +impl SchedulerCore + for RoundRobinCore { - type Sink = DispatchQueueSinkType; type StorageClient = SchedulerStorageClientType; + fn get_dispatch_queue_handle(&self) -> SharedDispatchQueueHandle { + self.dispatch_queue_reader.clone() + } + async fn run( self: Box, storage_client: Self::StorageClient, - sink: Self::Sink, reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver, id_issuer: TaskAssignmentIdIssuer, cancellation_token: CancellationToken, ) -> Result<(), SchedulerError> { + let Self { + config, + dispatch_queue_writer, + .. + } = *self; RoundRobin::new( SessionId::default(), storage_client, - sink, + dispatch_queue_writer, reschedule_queue_reader, id_issuer, cancellation_token, - self.config, + config, ) .run() .await @@ -172,17 +178,13 @@ impl JobTaskQueue { /// # Type Parameters /// /// * `SchedulerStorageClientType` - The storage client used to poll the inbound queue. -/// * `DispatchQueueSinkType` - The dispatch sink that task assignments are written to. /// /// # Note /// /// All member variables are marked `pub(super)` to allow the test module to inspect the internal /// states. -pub(super) struct RoundRobin< - SchedulerStorageClientType: SchedulerStorageClient + 'static, - DispatchQueueSinkType: DispatchQueueSink, -> { - pub(super) sink: DispatchQueueSinkType, +pub(super) struct RoundRobin { + pub(super) dispatch_queue_writer: DispatchQueueWriter, pub(super) cancellation_token: CancellationToken, pub(super) id_issuer: TaskAssignmentIdIssuer, pub(super) config: RoundRobinConfig, @@ -207,10 +209,8 @@ pub(super) struct RoundRobin< pub(super) reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver, } -impl< - SchedulerStorageClientType: SchedulerStorageClient + 'static, - DispatchQueueSinkType: DispatchQueueSink, -> RoundRobin +impl + RoundRobin { /// Factory function. /// @@ -222,7 +222,7 @@ impl< pub(super) fn new( storage_session_id: SessionId, storage_client: SchedulerStorageClientType, - sink: DispatchQueueSinkType, + dispatch_queue_writer: DispatchQueueWriter, reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver, id_issuer: TaskAssignmentIdIssuer, cancellation_token: CancellationToken, @@ -242,7 +242,7 @@ impl< let finalizing_job_queue = VecDeque::new(); let inbound_queue_reader = AsyncInboundQueueReader::new(storage_client); Self { - sink, + dispatch_queue_writer, cancellation_token, id_issuer, config, @@ -461,7 +461,7 @@ impl< /// /// * [`SchedulerError::InvalidSessionId`] if the polled session is older than the current /// session. - /// * Forwards [`DispatchQueueSink::bump_session_id`]'s return values on failure. + /// * Forwards [`DispatchQueueWriter::bump_session_id`]'s return values on failure. /// * Forwards [`Self::enqueue_commit_ready_entries`]'s return values on failure. /// * Forwards [`Self::enqueue_cleanup_ready_entries`]'s return values on failure. async fn ingest_inbound_entries( @@ -484,7 +484,9 @@ impl< ); self.storage_session_id = storage_session_id; self.clear(); - self.sink.bump_session_id(storage_session_id).await?; + self.dispatch_queue_writer + .bump_session_id(storage_session_id) + .await?; } // Load commit-ready tasks and cleanup-ready tasks first to avoid loading a job that is @@ -776,14 +778,14 @@ impl< /// /// * [`SchedulerError::Internal`] if the round-robin queue is inconsistent with the scheduler's /// job bookkeeping. - /// * Forwards [`DispatchQueueSink::enqueue`]'s return values on failure. + /// * Forwards [`DispatchQueueWriter::enqueue`]'s return values on failure. /// * Forwards [`Self::retire_active_job`]'s return values on failure. async fn make_schedule_decisions(&mut self) -> Result<(), SchedulerError> { let dispatch_slots = self .config .dispatch_queue_capacity .get() - .saturating_sub(self.sink.size()); + .saturating_sub(self.dispatch_queue_writer.size()); let mut remaining_dispatch_slots = dispatch_slots; 'fill_dispatch_queue: while remaining_dispatch_slots > 0 && !self.buffered_tasks.is_empty() { @@ -806,7 +808,7 @@ impl< else { continue; }; - self.sink + self.dispatch_queue_writer .enqueue(TaskAssignment { id: self.id_issuer.next(), job_id, @@ -827,7 +829,7 @@ impl< else { break; }; - self.sink + self.dispatch_queue_writer .enqueue(TaskAssignment { id: self.id_issuer.next(), job_id, @@ -847,7 +849,7 @@ impl< ))); }; if let Some(task_id) = job_entry.dequeue() { - self.sink + self.dispatch_queue_writer .enqueue(TaskAssignment { id: self.id_issuer.next(), job_id, diff --git a/components/spider-scheduler/src/core_impl/round_robin/tests.rs b/components/spider-scheduler/src/core_impl/round_robin/tests.rs index fdaacac06..9f62dee6b 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/tests.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/tests.rs @@ -24,7 +24,6 @@ use tokio_util::sync::CancellationToken; use super::RoundRobinConfig; use super::implementation::RoundRobin; -use crate::DispatchQueueSource; use crate::InboundEntry; use crate::SchedulerCore; use crate::SchedulerError; @@ -32,19 +31,28 @@ use crate::SchedulerStorageClient; use crate::StorageClientError; use crate::TaskAssignment; use crate::core::TaskAssignmentIdIssuer; +use crate::dispatch_queue::DispatchQueueHandle; use crate::dispatch_queue::DispatchQueueReader; -use crate::dispatch_queue::DispatchQueueWriter; +use crate::dispatch_queue::SharedDispatchQueueHandle; use crate::dispatch_queue::create_dispatch_queue; -/// The session used by tests that never bump the session. +/// The session used by tests that never bump the session. Must stay equal to the session +/// [`RoundRobinConfig::make_core`] seeds its dispatch queue with, so that the white-box helpers +/// below build a queue equivalent to the one the core owns. const DEFAULT_SESSION_ID: SessionId = 0; -/// The white-box scheduler under test, driven by manual ticks. -type TestScheduler = RoundRobin; - /// The maximum time to wait for expected assignments before failing a test. const DRAIN_DEADLINE: Duration = Duration::from_secs(5); +/// The white-box scheduler under test, driven by manual ticks. +type TestScheduler = RoundRobin; + +/// The join handle yielding a spawned scheduler's exit result. +type SchedulerJoinHandle = tokio::task::JoinHandle>; + +/// The sender a test pushes assignments onto as a lost execution manager's registry entry would. +type RescheduleQueueSender = tokio::sync::mpsc::UnboundedSender; + struct MockStorageInner { session_id: AtomicU64, ready_batches: Mutex)>>, @@ -295,17 +303,18 @@ fn make_reschedule_assignment( /// /// * The join handle yielding the scheduler's exit result. /// * The cancellation token that stops the scheduler. +/// * The handle over the dispatch queue the core owns, playing the execution manager's role. fn spawn_scheduler( config: RoundRobinConfig, storage_client: MockStorageClient, - sink: DispatchQueueWriter, ) -> ( - tokio::task::JoinHandle>, + SchedulerJoinHandle, CancellationToken, + SharedDispatchQueueHandle, ) { - let (handle, cancellation_token, _reschedule_queue_sender) = - spawn_scheduler_with_reschedule(config, storage_client, sink); - (handle, cancellation_token) + let (handle, cancellation_token, dispatch_queue_handle, _reschedule_queue_sender) = + spawn_scheduler_with_reschedule(config, storage_client); + (handle, cancellation_token, dispatch_queue_handle) } /// Spawns the scheduler's public run loop as a background task, exposing the reschedule-queue @@ -317,35 +326,42 @@ fn spawn_scheduler( /// /// * The join handle yielding the scheduler's exit result. /// * The cancellation token that stops the scheduler. -/// * The reschedule-queue sender a test uses to push assignments back as if a worker took them then -/// died. +/// * The handle over the dispatch queue the core owns, playing the execution manager's role. +/// * The reschedule-queue sender a test uses to push assignments back as a lost execution manager +/// would have returned them. fn spawn_scheduler_with_reschedule( config: RoundRobinConfig, storage_client: MockStorageClient, - sink: DispatchQueueWriter, ) -> ( - tokio::task::JoinHandle>, + SchedulerJoinHandle, CancellationToken, - tokio::sync::mpsc::UnboundedSender, + SharedDispatchQueueHandle, + RescheduleQueueSender, ) { - let core = Box::new(config.make_core()); + let core = Box::new(config.make_core::()); + let dispatch_queue_handle = core.get_dispatch_queue_handle(); let cancellation_token = CancellationToken::new(); let scheduler_token = cancellation_token.clone(); let (reschedule_queue_sender, reschedule_queue_reader) = tokio::sync::mpsc::unbounded_channel(); let handle = tokio::spawn(async move { core.run( storage_client, - sink, reschedule_queue_reader, TaskAssignmentIdIssuer::new(), scheduler_token, ) .await }); - (handle, cancellation_token, reschedule_queue_sender) + ( + handle, + cancellation_token, + dispatch_queue_handle, + reschedule_queue_sender, + ) } -/// Drains exactly `n` task assignments from the dispatch queue, playing the worker pool's role. +/// Drains exactly `n` task assignments from the dispatch queue, playing the execution manager's +/// role. /// /// # Returns /// @@ -356,8 +372,11 @@ fn spawn_scheduler_with_reschedule( /// Returns an error if: /// /// * Fewer than `n` assignments arrive within [`DRAIN_DEADLINE`]. -/// * Forwards [`DispatchQueueSource::dequeue`]'s return values on failure. -async fn drain_n(reader: &DispatchQueueReader, n: usize) -> anyhow::Result> { +/// * Forwards [`DispatchQueueHandle::dequeue`]'s return values on failure. +async fn drain_n( + dispatch_queue_handle: &dyn DispatchQueueHandle, + n: usize, +) -> anyhow::Result> { const DEQUEUE_WAIT: Duration = Duration::from_millis(100); let deadline = tokio::time::Instant::now() + DRAIN_DEADLINE; let mut assignments = Vec::with_capacity(n); @@ -368,7 +387,7 @@ async fn drain_n(reader: &DispatchQueueReader, n: usize) -> anyhow::Result anyhow::Result anyhow::Result<()> { +async fn assert_no_more_assignments( + dispatch_queue_handle: &dyn DispatchQueueHandle, +) -> anyhow::Result<()> { const OBSERVATION_WINDOW: Duration = Duration::from_secs(1); - let unexpected_assignment = reader.dequeue(OBSERVATION_WINDOW).await?; + let unexpected_assignment = dispatch_queue_handle.dequeue(OBSERVATION_WINDOW).await?; assert_eq!(unexpected_assignment, None); Ok(()) } @@ -488,40 +509,46 @@ fn assert_round_robin_property( /// # Returns /// -/// A white-box scheduler wired to the given storage client and sink, to be driven by manual -/// [`RoundRobin::tick`] calls. +/// A tuple containing: +/// +/// * A white-box scheduler wired to the given storage client, to be driven by manual +/// [`RoundRobin::tick`] calls. +/// * The reader of the dispatch queue the scheduler writes to. fn make_scheduler( config: RoundRobinConfig, storage_client: MockStorageClient, - sink: DispatchQueueWriter, -) -> TestScheduler { - make_scheduler_with_reschedule(config, storage_client, sink).0 +) -> (TestScheduler, DispatchQueueReader) { + let (scheduler, reader, _reschedule_queue_sender) = + make_scheduler_with_reschedule(config, storage_client); + (scheduler, reader) } /// # Returns /// -/// A white-box scheduler wired to the given storage client and sink, driven by manual -/// [`RoundRobin::tick`] calls, together with the reschedule-queue sender a test uses to inject -/// the assignments a lost execution manager would have returned. +/// A tuple containing: +/// +/// * A white-box scheduler wired to the given storage client, to be driven by manual +/// [`RoundRobin::tick`] calls. +/// * The reader of the dispatch queue the scheduler writes to. +/// * The reschedule-queue sender a test uses to inject the assignments a lost execution manager +/// would have returned. fn make_scheduler_with_reschedule( config: RoundRobinConfig, storage_client: MockStorageClient, - sink: DispatchQueueWriter, -) -> ( - TestScheduler, - tokio::sync::mpsc::UnboundedSender, -) { +) -> (TestScheduler, DispatchQueueReader, RescheduleQueueSender) { + let (writer, reader) = + create_dispatch_queue(config.dispatch_queue_capacity.get(), DEFAULT_SESSION_ID); let (reschedule_queue_sender, reschedule_queue_reader) = tokio::sync::mpsc::unbounded_channel(); let scheduler = RoundRobin::new( DEFAULT_SESSION_ID, storage_client, - sink, + writer, reschedule_queue_reader, TaskAssignmentIdIssuer::new(), CancellationToken::new(), config, ); - (scheduler, reschedule_queue_sender) + (scheduler, reader, reschedule_queue_sender) } /// Ticks the scheduler until `predicate` holds on its state. @@ -561,7 +588,7 @@ async fn tick_until( /// /// * Fewer than `n` assignments arrive within [`DRAIN_DEADLINE`]. /// * Forwards [`RoundRobin::tick`]'s return values on failure. -/// * Forwards [`DispatchQueueSource::dequeue`]'s return values on failure. +/// * Forwards [`DispatchQueueHandle::dequeue`]'s return values on failure. async fn tick_and_drain_n( scheduler: &mut TestScheduler, reader: &DispatchQueueReader, @@ -592,7 +619,7 @@ async fn tick_and_drain_n( /// Returns an error if: /// /// * Forwards [`RoundRobin::tick`]'s return values on failure. -/// * Forwards [`DispatchQueueSource::dequeue`]'s return values on failure. +/// * Forwards [`DispatchQueueHandle::dequeue`]'s return values on failure. /// /// # Panics /// @@ -665,11 +692,9 @@ async fn assert_finalizing_ready_drops_jobs(finalizing_task_id: TaskId) -> anyho make_ready_batch(&jobs, TASKS_PER_JOB, 0), ); - let (writer, reader) = create_dispatch_queue(DISPATCH_QUEUE_CAPACITY, DEFAULT_SESSION_ID); - let mut scheduler = make_scheduler( + let (mut scheduler, reader) = make_scheduler( make_config(ACTIVE_JOB_QUEUE_CAPACITY, DISPATCH_QUEUE_CAPACITY), storage_client.clone(), - writer, ); // Step 1: ingest the ready batch. The ingesting tick also dispatches exactly two assignments @@ -843,12 +868,12 @@ async fn single_capacity_pool_schedules_jobs_serially() -> anyhow::Result<()> { make_ready_batch(&jobs, TASKS_PER_JOB, DUP_EVERY), ); - let (writer, reader) = create_dispatch_queue(DISPATCH_QUEUE_CAPACITY, DEFAULT_SESSION_ID); let config = make_config(1, DISPATCH_QUEUE_CAPACITY); - let (scheduler_handle, cancellation_token) = spawn_scheduler(config, storage_client, writer); + let (scheduler_handle, cancellation_token, dispatch_queue_handle) = + spawn_scheduler(config, storage_client); - let assignments = drain_n(&reader, NUM_JOBS * TASKS_PER_JOB).await?; - assert_no_more_assignments(&reader).await?; + let assignments = drain_n(&*dispatch_queue_handle, NUM_JOBS * TASKS_PER_JOB).await?; + assert_no_more_assignments(&*dispatch_queue_handle).await?; // With an active job pool of capacity 1, round-robin degenerates to serial job FIFO: the // rotation holds a single job at a time, so each job's tasks dispatch as one consecutive @@ -876,12 +901,12 @@ async fn active_jobs_dispatch_in_round_robin_order() -> anyhow::Result<()> { make_ready_batch(&jobs, TASKS_PER_JOB, DUP_EVERY), ); - let (writer, reader) = create_dispatch_queue(DISPATCH_QUEUE_CAPACITY, DEFAULT_SESSION_ID); let config = make_config(NUM_JOBS, DISPATCH_QUEUE_CAPACITY); - let (scheduler_handle, cancellation_token) = spawn_scheduler(config, storage_client, writer); + let (scheduler_handle, cancellation_token, dispatch_queue_handle) = + spawn_scheduler(config, storage_client); - let assignments = drain_n(&reader, NUM_JOBS * TASKS_PER_JOB).await?; - assert_no_more_assignments(&reader).await?; + let assignments = drain_n(&*dispatch_queue_handle, NUM_JOBS * TASKS_PER_JOB).await?; + assert_no_more_assignments(&*dispatch_queue_handle).await?; // All 10 jobs fit into the active job pool, so no job ever pends and dispatch follows the // strict rotation: task 0 of every job in batch order, then task 1 of every job, and so on. The @@ -909,12 +934,12 @@ async fn pending_jobs_promote_and_schedule_round_robin() -> anyhow::Result<()> { make_ready_batch(&jobs, TASKS_PER_JOB, DUP_EVERY), ); - let (writer, reader) = create_dispatch_queue(DISPATCH_QUEUE_CAPACITY, DEFAULT_SESSION_ID); let config = make_config(ACTIVE_JOB_QUEUE_CAPACITY, DISPATCH_QUEUE_CAPACITY); - let (scheduler_handle, cancellation_token) = spawn_scheduler(config, storage_client, writer); + let (scheduler_handle, cancellation_token, dispatch_queue_handle) = + spawn_scheduler(config, storage_client); - let assignments = drain_n(&reader, NUM_JOBS * TASKS_PER_JOB).await?; - assert_no_more_assignments(&reader).await?; + let assignments = drain_n(&*dispatch_queue_handle, NUM_JOBS * TASKS_PER_JOB).await?; + assert_no_more_assignments(&*dispatch_queue_handle).await?; let (active_jobs, pending_jobs) = jobs.split_at(ACTIVE_JOB_QUEUE_CAPACITY); let (phase1, phase2) = assignments.split_at(ACTIVE_JOB_QUEUE_CAPACITY * TASKS_PER_JOB); @@ -960,14 +985,14 @@ async fn commit_drains_each_cycle_cleanup_dispatches_once() -> anyhow::Result<() make_finalizing_batch(&cleanup_ready_jobs, TaskId::Cleanup), ); - let (writer, reader) = create_dispatch_queue(DISPATCH_QUEUE_CAPACITY, DEFAULT_SESSION_ID); let config = make_config(NUM_ACTIVE_JOBS, DISPATCH_QUEUE_CAPACITY); - let (scheduler_handle, cancellation_token) = spawn_scheduler(config, storage_client, writer); + let (scheduler_handle, cancellation_token, dispatch_queue_handle) = + spawn_scheduler(config, storage_client); let num_assignments = NUM_ACTIVE_JOBS * TASKS_PER_JOB + NUM_COMMIT_READY_JOBS + NUM_CLEANUP_READY_JOBS; - let assignments = drain_n(&reader, num_assignments).await?; - assert_no_more_assignments(&reader).await?; + let assignments = drain_n(&*dispatch_queue_handle, num_assignments).await?; + assert_no_more_assignments(&*dispatch_queue_handle).await?; // The rotation is [commit lane, cleanup lane, active jobs...]. The commit lane drains up to // `active_job_queue_capacity` (== NUM_ACTIVE_JOBS) jobs per visit, so each cycle dispatches a @@ -1026,11 +1051,9 @@ async fn session_bump_clears_buffered_tasks() -> anyhow::Result<()> { make_ready_batch(&old_jobs, TASKS_PER_JOB, 0), ); - let (writer, reader) = create_dispatch_queue(DISPATCH_QUEUE_CAPACITY, DEFAULT_SESSION_ID); - let mut scheduler = make_scheduler( + let (mut scheduler, reader) = make_scheduler( make_config(ACTIVE_JOB_QUEUE_CAPACITY, DISPATCH_QUEUE_CAPACITY), storage_client.clone(), - writer, ); // Step 1: ingest the old-session batch. The ingesting tick dispatches enough assignments to @@ -1105,11 +1128,9 @@ async fn reschedule_ready_task_redispatches_with_current_session() -> anyhow::Re let job_a = jobs[0]; let storage_client = MockStorageClient::new(DEFAULT_SESSION_ID); - let (writer, reader) = create_dispatch_queue(DISPATCH_QUEUE_CAPACITY, DEFAULT_SESSION_ID); - let (mut scheduler, reschedule_queue_sender) = make_scheduler_with_reschedule( + let (mut scheduler, reader, reschedule_queue_sender) = make_scheduler_with_reschedule( make_config(ACTIVE_JOB_QUEUE_CAPACITY, DISPATCH_QUEUE_CAPACITY), storage_client, - writer, ); let injected = make_reschedule_assignment(job_a, TaskId::Index(0), DEFAULT_SESSION_ID, 0); @@ -1138,11 +1159,9 @@ async fn reschedule_drops_stale_session_assignment() -> anyhow::Result<()> { let storage_client = MockStorageClient::new(DEFAULT_SESSION_ID); storage_client.set_session(NEW_SESSION_ID); - let (writer, reader) = create_dispatch_queue(DISPATCH_QUEUE_CAPACITY, DEFAULT_SESSION_ID); - let (mut scheduler, reschedule_queue_sender) = make_scheduler_with_reschedule( + let (mut scheduler, reader, reschedule_queue_sender) = make_scheduler_with_reschedule( make_config(ACTIVE_JOB_QUEUE_CAPACITY, DISPATCH_QUEUE_CAPACITY), storage_client, - writer, ); // An empty poll under the higher session bumps the scheduler's session past the stale @@ -1185,11 +1204,9 @@ async fn reschedule_dedups_against_inbound_ready_task() -> anyhow::Result<()> { let storage_client = MockStorageClient::new(DEFAULT_SESSION_ID); storage_client.push_ready_batch(DEFAULT_SESSION_ID, make_ready_batch(&jobs, 1, 0)); - let (writer, reader) = create_dispatch_queue(DISPATCH_QUEUE_CAPACITY, DEFAULT_SESSION_ID); - let (mut scheduler, reschedule_queue_sender) = make_scheduler_with_reschedule( + let (mut scheduler, reader, reschedule_queue_sender) = make_scheduler_with_reschedule( make_config(ACTIVE_JOB_QUEUE_CAPACITY, DISPATCH_QUEUE_CAPACITY), storage_client, - writer, ); tick_until(&mut scheduler, |scheduler| { @@ -1233,12 +1250,11 @@ async fn randomly_rescheduled_assignments_are_eventually_redispatched() -> anyho make_ready_batch(&jobs, TASKS_PER_JOB, 0), ); - let (writer, reader) = create_dispatch_queue(DISPATCH_QUEUE_CAPACITY, DEFAULT_SESSION_ID); - let (scheduler_handle, cancellation_token, reschedule_sender) = spawn_scheduler_with_reschedule( - make_config(NUM_JOBS, DISPATCH_QUEUE_CAPACITY), - storage_client, - writer, - ); + let (scheduler_handle, cancellation_token, dispatch_queue_handle, reschedule_sender) = + spawn_scheduler_with_reschedule( + make_config(NUM_JOBS, DISPATCH_QUEUE_CAPACITY), + storage_client, + ); let total = NUM_JOBS * TASKS_PER_JOB; let mut completed: HashSet<(JobId, TaskId)> = HashSet::new(); @@ -1254,7 +1270,10 @@ async fn randomly_rescheduled_assignments_are_eventually_redispatched() -> anyho completed.len(), ); } - let Some(assignment) = reader.dequeue(Duration::from_millis(100)).await? else { + let Some(assignment) = dispatch_queue_handle + .dequeue(Duration::from_millis(100)) + .await? + else { continue; }; assert_eq!(assignment.session_id, DEFAULT_SESSION_ID); diff --git a/components/spider-scheduler/src/dispatch_queue.rs b/components/spider-scheduler/src/dispatch_queue.rs index 23ccab176..fc300bc6d 100644 --- a/components/spider-scheduler/src/dispatch_queue.rs +++ b/components/spider-scheduler/src/dispatch_queue.rs @@ -11,46 +11,10 @@ use tokio::sync::RwLock; use crate::error::SchedulerError; use crate::types::TaskAssignment; -/// The writer side of the dispatching queue used by the scheduler core. +/// The access point to a scheduler core's dispatching queue, used by the execution-manager-facing +/// service to drain the queue. #[async_trait] -pub trait DispatchQueueSink: Send + Sync + Clone { - /// Enqueues a task assignment for execution managers to consume. - /// - /// # Parameters - /// - /// * `assignment` - The task assignment to enqueue. - /// - /// # Errors - /// - /// Returns an error if: - /// - /// * [`SchedulerError::DispatchQueueClosed`] if the dispatching queue is closed. - async fn enqueue(&self, assignment: TaskAssignment) -> Result<(), SchedulerError>; - - /// Bumps the session ID and invalidates all queued task assignments. - /// - /// # Parameters - /// - /// * `new_session_id` - The new session ID. Must be greater than the current session ID. - /// - /// # Errors - /// - /// Returns an error if: - /// - /// * [`SchedulerError::DispatchQueueClosed`] if the dispatching queue is closed. - /// * [`SchedulerError::InvalidSessionId`] if the new session ID is not greater than the current - /// session ID. - async fn bump_session_id(&self, new_session_id: SessionId) -> Result<(), SchedulerError>; - - /// # Returns - /// - /// The current size of the dispatch queue. - fn size(&self) -> usize; -} - -/// The reader side of the dispatching queue, drained by the execution-manager-facing service. -#[async_trait] -pub trait DispatchQueueSource: Send + Sync + Clone { +pub trait DispatchQueueHandle: Send + Sync { /// Dequeues the next task assignment for an execution manager to execute. /// /// # Parameters @@ -70,8 +34,11 @@ pub trait DispatchQueueSource: Send + Sync + Clone { async fn dequeue(&self, wait_time: Duration) -> Result, SchedulerError>; } -/// A cloneable writer handle for the dispatching queue, implementing [`DispatchQueueSink`] using -/// an async channel. +/// The shared handle over a dispatching queue that a scheduler core hands to the +/// execution-manager-facing service. +pub type SharedDispatchQueueHandle = Arc; + +/// A cloneable writer handle for the dispatching queue, backed by an async channel. /// /// # NOTE /// @@ -83,9 +50,15 @@ pub struct DispatchQueueWriter { inner: Arc, } -#[async_trait] -impl DispatchQueueSink for DispatchQueueWriter { - async fn enqueue(&self, assignment: TaskAssignment) -> Result<(), SchedulerError> { +impl DispatchQueueWriter { + /// Enqueues a task assignment for execution managers to consume. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`SchedulerError::DispatchQueueClosed`] if the dispatching queue is closed. + pub async fn enqueue(&self, assignment: TaskAssignment) -> Result<(), SchedulerError> { self.inner .assignment_sender .send(assignment) @@ -93,7 +66,16 @@ impl DispatchQueueSink for DispatchQueueWriter { .map_err(|_| SchedulerError::DispatchQueueClosed) } - async fn bump_session_id(&self, new_session_id: SessionId) -> Result<(), SchedulerError> { + /// Bumps the session ID and invalidates all queued task assignments. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`SchedulerError::DispatchQueueClosed`] if the dispatching queue is closed. + /// * [`SchedulerError::InvalidSessionId`] if the new session ID is not greater than the current + /// session ID. + pub async fn bump_session_id(&self, new_session_id: SessionId) -> Result<(), SchedulerError> { let mut session_id_guard = self.inner.session_id.write().await; if new_session_id <= *session_id_guard { return Err(SchedulerError::InvalidSessionId(new_session_id)); @@ -108,26 +90,28 @@ impl DispatchQueueSink for DispatchQueueWriter { Ok(()) } - fn size(&self) -> usize { + #[must_use] + pub fn size(&self) -> usize { self.inner.assignment_sender.len() } } -/// A cloneable reader handle for the dispatching queue, implementing [`DispatchQueueSource`] using +/// A cloneable reader handle for the dispatching queue, implementing [`DispatchQueueHandle`] using /// an async channel. #[derive(Clone)] pub struct DispatchQueueReader { - inner: Arc, + session_id: Arc>, + assignment_receiver: async_channel::Receiver, } #[async_trait] -impl DispatchQueueSource for DispatchQueueReader { +impl DispatchQueueHandle for DispatchQueueReader { async fn dequeue(&self, wait_time: Duration) -> Result, SchedulerError> { // Lock session ID for the entire duration of the dequeue operation to exclude any // `bump_session_id` operations. - let _session_id_guard = self.inner.session_id.read().await; + let _session_id_guard = self.session_id.read().await; - if let Ok(assignment) = self.inner.assignment_receiver.try_recv() { + if let Ok(assignment) = self.assignment_receiver.try_recv() { return Ok(Some(assignment)); } @@ -135,7 +119,7 @@ impl DispatchQueueSource for DispatchQueueReader { return Ok(None); } - match tokio::time::timeout(wait_time, self.inner.assignment_receiver.recv()).await { + match tokio::time::timeout(wait_time, self.assignment_receiver.recv()).await { Ok(Ok(assignment)) => Ok(Some(assignment)), Ok(Err(_)) => Err(SchedulerError::DispatchQueueClosed), Err(_) => Ok(None), @@ -163,16 +147,13 @@ pub fn create_dispatch_queue( assignment_sender, assignment_receiver: assignment_receiver.clone(), }); - let reader_inner = Arc::new(DispatchQueueReaderInner { - session_id, - assignment_receiver, - }); ( DispatchQueueWriter { inner: writer_inner, }, DispatchQueueReader { - inner: reader_inner, + session_id, + assignment_receiver, }, ) } @@ -183,11 +164,6 @@ struct DispatchQueueWriterInner { assignment_receiver: async_channel::Receiver, } -struct DispatchQueueReaderInner { - session_id: Arc>, - assignment_receiver: async_channel::Receiver, -} - #[cfg(test)] mod tests { use std::collections::HashMap; diff --git a/components/spider-scheduler/src/grpc.rs b/components/spider-scheduler/src/grpc.rs index 5cbb196cc..9955c6f1d 100644 --- a/components/spider-scheduler/src/grpc.rs +++ b/components/spider-scheduler/src/grpc.rs @@ -15,37 +15,26 @@ use tonic::Request; use tonic::Response; use tonic::Status; -use crate::dispatch_queue::DispatchQueueSource; use crate::error::SchedulerError; use crate::error::SchedulerServiceError; use crate::execution_manager_registry::ExecutionManagerRegistryError; use crate::service::SchedulerServiceState; /// gRPC adapter over a [`SchedulerServiceState`]. -/// -/// # Type Parameters -/// -/// * `DispatchQueueSourceType` - The reader side of the dispatching queue the underlying service -/// drains. #[derive(Clone)] -pub struct GrpcSchedulerService { - inner: SchedulerServiceState, +pub struct GrpcSchedulerService { + inner: SchedulerServiceState, cancellation_token: CancellationToken, } -impl - GrpcSchedulerService -{ +impl GrpcSchedulerService { /// Factory function. /// /// # Returns /// /// A new [`GrpcSchedulerService`] wrapping [`SchedulerServiceState`]. #[must_use] - pub const fn new( - inner: SchedulerServiceState, - cancellation_token: CancellationToken, - ) -> Self { + pub const fn new(inner: SchedulerServiceState, cancellation_token: CancellationToken) -> Self { Self { inner, cancellation_token, @@ -117,9 +106,7 @@ impl /// All possible errors that can occur during scheduling can be found in /// [`GrpcSchedulerService::service_error_handler`]. #[async_trait] -impl SchedulerService - for GrpcSchedulerService -{ +impl SchedulerService for GrpcSchedulerService { async fn next_task( &self, request: Request, diff --git a/components/spider-scheduler/src/lib.rs b/components/spider-scheduler/src/lib.rs index a39e4d7b6..7d692e972 100644 --- a/components/spider-scheduler/src/lib.rs +++ b/components/spider-scheduler/src/lib.rs @@ -7,24 +7,26 @@ //! //! The crate defines three trait seams wired into a single pipeline — a storage client that polls //! the inbound queue, a core that makes serial decisions, and a dispatching queue that fans those -//! decisions out to execution managers: +//! decisions out to execution managers. The core owns and constructs its own dispatching queue, +//! and exposes only a handle to the service: //! //! ```text //! storage ── authoritative inbound queue (owned by the storage layer, not this crate) //! │ //! │ poll_ready / poll_commit_ready / poll_cleanup_ready (SchedulerStorageClient) //! ▼ -//! ┌───────────────────┐ -//! │ SchedulerCore │ serial loop: poll → decide → enqueue -//! └───────────────────┘ -//! │ -//! │ enqueue (DispatchQueueSink — writer side) -//! ▼ -//! ┌───────────────────┐ -//! │ dispatch queue │ bounded SPMC; a full queue back-pressures the core -//! └───────────────────┘ +//! ┌───────────────────────────────────────────────────┐ +//! │ SchedulerCore │ +//! │ serial loop: poll → decide → enqueue │ +//! │ │ +//! │ ┌───────────────────────────────────────────┐ │ +//! │ │ dispatch queue │ │ +//! │ │ owned and constructed by the core; │ │ +//! │ │ a full queue back-pressures the core │ │ +//! │ └───────────────────────────────────────────┘ │ +//! └───────────────────────────────────────────────────┘ //! │ -//! │ dequeue (DispatchQueueSource — reader side) +//! │ dequeue (DispatchQueueHandle) //! ▼ //! ┌───────────────────┐ //! │ scheduler service │ ──▶ execution managers (concurrent fan-out) @@ -46,10 +48,10 @@ pub mod types; pub use crate::config::SchedulerConfig; pub use crate::config::ServerConfig; pub use crate::core::SchedulerCore; +pub use crate::dispatch_queue::DispatchQueueHandle; pub use crate::dispatch_queue::DispatchQueueReader; -pub use crate::dispatch_queue::DispatchQueueSink; -pub use crate::dispatch_queue::DispatchQueueSource; pub use crate::dispatch_queue::DispatchQueueWriter; +pub use crate::dispatch_queue::SharedDispatchQueueHandle; pub use crate::error::SchedulerError; pub use crate::error::SchedulerRuntimeError; pub use crate::error::SchedulerServiceError; diff --git a/components/spider-scheduler/src/runtime.rs b/components/spider-scheduler/src/runtime.rs index ec77077a7..9d6e2e286 100644 --- a/components/spider-scheduler/src/runtime.rs +++ b/components/spider-scheduler/src/runtime.rs @@ -1,23 +1,19 @@ //! The scheduler runtime. //! -//! This module registers the scheduler with the storage service, wires the scheduler core to a -//! freshly created dispatch queue, hands the core the reschedule queue reader, and spawns the -//! core's scheduling loop as a background coroutine alongside the execution manager registry. The -//! resulting [`Runtime`] owns the spawned coroutine and is responsible for cancelling and joining -//! it on shutdown. +//! This module registers the scheduler with the storage service, builds the scheduler core and +//! wires the execution-manager-facing service to the handle of the dispatch queue the core owns, +//! hands the core the reschedule queue reader, and spawns the core's scheduling loop as a +//! background coroutine alongside the execution manager registry. The resulting [`Runtime`] owns +//! the spawned coroutine and is responsible for cancelling and joining it on shutdown. use std::time::Duration; use serde::Deserialize; -use spider_core::types::id::SessionId; use spider_utils::config::EndpointConfig; use tokio_util::sync::CancellationToken; use crate::config::SchedulerConfig; use crate::core::TaskAssignmentIdIssuer; -use crate::dispatch_queue::DispatchQueueReader; -use crate::dispatch_queue::DispatchQueueWriter; -use crate::dispatch_queue::create_dispatch_queue; use crate::error::SchedulerError; use crate::error::SchedulerRuntimeError; use crate::execution_manager_registry::ExecutionManagerRegistry; @@ -88,9 +84,10 @@ impl Runtime { /// Creates a scheduler runtime from the given configuration and storage client. /// -/// Registers this scheduler with the storage service, wires the scheduler core to a freshly created -/// dispatch queue, hands the core the reschedule queue reader, and starts the core's scheduling -/// loop as a background coroutine. +/// Registers this scheduler with the storage service, builds the scheduler core and wires the +/// execution-manager-facing service to the handle of the dispatch queue the core owns, hands the +/// core the reschedule queue reader, and starts the core's scheduling loop as a background +/// coroutine. /// /// # Type Parameters /// @@ -101,7 +98,7 @@ impl Runtime { /// A tuple on success, containing: /// /// * The newly created runtime instance. -/// * The execution-manager-facing scheduler service, built over the dispatch queue reader. +/// * The execution-manager-facing scheduler service, built over the core's dispatch queue handle. /// * The runtime's cancellation token for cancelling the runtime on error. /// /// # Errors @@ -112,14 +109,7 @@ impl Runtime { pub async fn create_runtime( config: RuntimeConfig, storage_client: SchedulerStorageClientType, -) -> Result< - ( - Runtime, - SchedulerServiceState, - CancellationToken, - ), - SchedulerRuntimeError, -> { +) -> Result<(Runtime, SchedulerServiceState, CancellationToken), SchedulerRuntimeError> { let RuntimeConfig { scheduler: scheduler_config, em_registry: execution_manager_registry_config, @@ -132,7 +122,6 @@ pub async fn create_runtime(); + let core = scheduler_config.make_core::(); + let dispatch_queue_handle = core.get_dispatch_queue_handle(); + let service = SchedulerServiceState::new(dispatch_queue_handle, registry, scheduler_id); let core_cancellation_token = cancellation_token.clone(); let core_join_handle = tokio::spawn(async move { core.run( storage_client, - dispatch_queue_writer, reschedule_queue_receiver, TaskAssignmentIdIssuer::new(), core_cancellation_token.clone(), @@ -195,6 +182,7 @@ mod tests { use spider_core::job::JobState; use spider_core::types::id::JobId; use spider_core::types::id::SchedulerId; + use spider_core::types::id::SessionId; use spider_utils::config::Host; use super::*; diff --git a/components/spider-scheduler/src/service.rs b/components/spider-scheduler/src/service.rs index 36117b6d4..582e61950 100644 --- a/components/spider-scheduler/src/service.rs +++ b/components/spider-scheduler/src/service.rs @@ -3,8 +3,9 @@ //! [`SchedulerServiceState`] is the domain layer behind the scheduler gRPC service. It serves //! execution managers by draining task assignments from the dispatch queue and bookkeeping them in //! the [`ExecutionManagerRegistry`]: assignment, completion, heartbeat, and shutdown. The service -//! is generic over its dispatch source, so the runtime can drive it with the real dispatch queue in -//! production or a mock in tests, while the [`ExecutionManagerRegistry`] is shared by value. +//! holds its dispatch queue handle as a trait object, so the runtime can drive it with the handle +//! of whichever dispatch structure the scheduler core owns in production, or with a mock in tests, +//! while the [`ExecutionManagerRegistry`] is shared by value. use std::sync::Arc; use std::time::Duration; @@ -13,24 +14,18 @@ use spider_core::types::id::ExecutionManagerId; use spider_core::types::id::SchedulerId; use spider_core::types::scheduler::TaskAssignmentRecord; -use crate::dispatch_queue::DispatchQueueSource; +use crate::dispatch_queue::SharedDispatchQueueHandle; use crate::error::SchedulerServiceError; use crate::execution_manager_registry::ExecutionManagerRegistry; use crate::types::TaskAssignment; /// The execution-manager-facing scheduler service. -/// -/// # Type Parameters -/// -/// * `DispatchQueueSourceType` - The reader side of the dispatching queue the service drains. #[derive(Clone)] -pub struct SchedulerServiceState { - inner: Arc>, +pub struct SchedulerServiceState { + inner: Arc, } -impl - SchedulerServiceState -{ +impl SchedulerServiceState { /// Factory function. /// /// # Returns @@ -38,13 +33,13 @@ impl /// A newly constructed [`SchedulerServiceState`]. #[must_use] pub fn new( - dispatch_source: DispatchQueueSourceType, + dispatch_queue_handle: SharedDispatchQueueHandle, registry: ExecutionManagerRegistry, scheduler_id: SchedulerId, ) -> Self { Self { inner: Arc::new(SchedulerServiceStateInner { - dispatch_source, + dispatch_queue_handle, registry, scheduler_id, }), @@ -78,7 +73,9 @@ impl /// /// Returns an error if: /// - /// * Forwards [`DispatchQueueSource::dequeue`]'s return values on failure. + /// * Forwards [`DispatchQueueHandle::dequeue`]'s return values on failure. + /// + /// [`DispatchQueueHandle::dequeue`]: crate::dispatch_queue::DispatchQueueHandle::dequeue pub async fn next_task( &self, em_id: ExecutionManagerId, @@ -95,7 +92,7 @@ impl prev, )); } - match self.inner.dispatch_source.dequeue(wait_time).await? { + match self.inner.dispatch_queue_handle.dequeue(wait_time).await? { None => { tracing::info!( scheduler_id = % self.scheduler_id(), @@ -204,12 +201,8 @@ impl } /// The shared inner state of [`SchedulerServiceState`]. -/// -/// # Type Parameters -/// -/// * `DispatchQueueSourceType` - The reader side of the dispatching queue the service drains. -struct SchedulerServiceStateInner { - dispatch_source: DispatchQueueSourceType, +struct SchedulerServiceStateInner { + dispatch_queue_handle: SharedDispatchQueueHandle, registry: ExecutionManagerRegistry, scheduler_id: SchedulerId, } @@ -237,13 +230,13 @@ mod tests { use tokio_util::sync::CancellationToken; use super::SchedulerServiceState; - use crate::dispatch_queue::DispatchQueueSource; + use crate::dispatch_queue::DispatchQueueHandle; use crate::error::SchedulerError; use crate::execution_manager_registry::ExecutionManagerRegistry; use crate::execution_manager_registry::ExecutionManagerRegistryConfig; use crate::types::TaskAssignment; - /// The storage session the mock dispatch source pairs with every assignment it returns. + /// The storage session the mock dispatch queue handle pairs with every assignment it returns. const SESSION_ID: SessionId = 7; /// The scheduler identifier the service stamps onto assignments. @@ -255,31 +248,30 @@ mod tests { /// The maximum time to wait for a rescheduled assignment before failing a test. const RESCHEDULE_TIMEOUT: Duration = Duration::from_secs(2); - /// A [`DispatchQueueSource`] mock backed by a shared counter. + /// A [`DispatchQueueHandle`] mock backed by a counter. /// - /// Each [`DispatchQueueSource::dequeue`] claims one slot from the counter; while it is positive + /// Each [`DispatchQueueHandle::dequeue`] claims one slot from the counter; while it is positive /// it returns a freshly minted task assignment, and once it reaches zero it returns [`None`] /// forever. Using a counter (rather than a canned list of assignments) lets the tests assert on /// dequeue behavior without coupling to recorded argument vectors. - #[derive(Clone)] - struct CounterDispatchSource { - remaining: Arc, + struct CounterDispatchQueueHandle { + remaining: AtomicUsize, } - impl CounterDispatchSource { + impl CounterDispatchQueueHandle { /// # Returns /// - /// A new [`CounterDispatchSource`] that hands out `remaining` assignments before reporting - /// an empty queue. - fn new(remaining: usize) -> Self { + /// A new [`CounterDispatchQueueHandle`] that hands out `remaining` assignments before + /// reporting an empty queue. + const fn new(remaining: usize) -> Self { Self { - remaining: Arc::new(AtomicUsize::new(remaining)), + remaining: AtomicUsize::new(remaining), } } } #[async_trait] - impl DispatchQueueSource for CounterDispatchSource { + impl DispatchQueueHandle for CounterDispatchQueueHandle { async fn dequeue( &self, _wait_time: Duration, @@ -363,20 +355,20 @@ mod tests { /// /// A tuple containing: /// - /// * A [`SchedulerServiceState`] over a [`CounterDispatchSource`] with `remaining` assignments, - /// backed by a fresh test registry. + /// * A [`SchedulerServiceState`] over a [`CounterDispatchQueueHandle`] with `remaining` + /// assignments, backed by a fresh test registry. /// * The receiver end of the registry's re-schedule queue. /// * The registry-level cancellation token. fn build_service( remaining: usize, ) -> ( - SchedulerServiceState, + SchedulerServiceState, UnboundedReceiver, CancellationToken, ) { let (registry, reschedule_queue_receiver, cancellation_token) = build_registry(); let service = SchedulerServiceState::new( - CounterDispatchSource::new(remaining), + Arc::new(CounterDispatchQueueHandle::new(remaining)), registry, SchedulerId::from(SCHEDULER_ID), );