From 0f2692818c9cf992aa318ca48bdd9c16b5d14f30 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Sun, 16 Aug 2026 22:28:36 -0400 Subject: [PATCH 1/4] refactor(spider-scheduler)!: Move dispatch queue ownership into the scheduler core: * Make the core own and construct its dispatch queue instead of receiving a pre-built sink as a parameter to `run`, so that a core is free to own a dispatch structure that is not a single channel. `SchedulerCore` loses its `Sink` associated type and gains `get_dispatch_queue_source`, and the runtime now builds the core first and takes the read handle from it. * Drop `Clone` from `DispatchQueueSource`'s supertraits so the trait is object-safe, and flatten `DispatchQueueReader` so the single `Arc` handed to the service is the only indirection. * Delete the `DispatchQueueSink` trait, whose sole implementor became a concrete type once the core owned its writer, moving `enqueue`, `bump_session_id` and `size` to an inherent `impl DispatchQueueWriter`. * Drop the now-unused `SchedulerConfig::dispatch_queue_capacity`, and the `DispatchQueueSourceType` type parameter from `SchedulerServiceState` and `GrpcSchedulerService`. --- components/spider-scheduler/src/config.rs | 25 +-- components/spider-scheduler/src/core.rs | 24 +-- .../core_impl/round_robin/implementation.rs | 84 ++++----- .../src/core_impl/round_robin/tests.rs | 172 ++++++++++-------- .../spider-scheduler/src/dispatch_queue.rs | 92 ++++------ components/spider-scheduler/src/grpc.rs | 23 +-- components/spider-scheduler/src/lib.rs | 25 +-- components/spider-scheduler/src/runtime.rs | 42 ++--- components/spider-scheduler/src/service.rs | 40 ++-- 9 files changed, 235 insertions(+), 292 deletions(-) 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..194195c11 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::DispatchQueueSource; 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 read handle over the core's dispatch structure, for the execution-manager-facing service + /// to drain. + fn get_dispatch_queue_source(&self) -> Arc; + /// 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..ca5453d44 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::DispatchQueueSource; +use crate::dispatch_queue::DispatchQueueWriter; +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_source: 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_source: Arc, + _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_source(&self) -> Arc { + Arc::clone(&self.dispatch_queue_source) + } + 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..8ee558925 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; @@ -33,18 +32,26 @@ use crate::StorageClientError; use crate::TaskAssignment; use crate::core::TaskAssignmentIdIssuer; use crate::dispatch_queue::DispatchQueueReader; -use crate::dispatch_queue::DispatchQueueWriter; +use crate::dispatch_queue::DispatchQueueSource; 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 +302,18 @@ fn make_reschedule_assignment( /// /// * The join handle yielding the scheduler's exit result. /// * The cancellation token that stops the scheduler. +/// * The read 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, + Arc, ) { - let (handle, cancellation_token, _reschedule_queue_sender) = - spawn_scheduler_with_reschedule(config, storage_client, sink); - (handle, cancellation_token) + let (handle, cancellation_token, dispatch_queue_source, _reschedule_queue_sender) = + spawn_scheduler_with_reschedule(config, storage_client); + (handle, cancellation_token, dispatch_queue_source) } /// Spawns the scheduler's public run loop as a background task, exposing the reschedule-queue @@ -317,35 +325,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 read 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, + Arc, + RescheduleQueueSender, ) { - let core = Box::new(config.make_core()); + let core = Box::new(config.make_core::()); + let dispatch_queue_source = core.get_dispatch_queue_source(); 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_source, + 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 /// @@ -357,7 +372,10 @@ fn spawn_scheduler_with_reschedule( /// /// * 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> { +async fn drain_n( + dispatch_queue_source: &dyn DispatchQueueSource, + 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 +386,7 @@ async fn drain_n(reader: &DispatchQueueReader, n: usize) -> anyhow::Result anyhow::Result anyhow::Result<()> { +async fn assert_no_more_assignments( + dispatch_queue_source: &dyn DispatchQueueSource, +) -> anyhow::Result<()> { const OBSERVATION_WINDOW: Duration = Duration::from_secs(1); - let unexpected_assignment = reader.dequeue(OBSERVATION_WINDOW).await?; + let unexpected_assignment = dispatch_queue_source.dequeue(OBSERVATION_WINDOW).await?; assert_eq!(unexpected_assignment, None); Ok(()) } @@ -488,40 +508,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. @@ -665,11 +691,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 +867,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_source) = + 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_source, NUM_JOBS * TASKS_PER_JOB).await?; + assert_no_more_assignments(&*dispatch_queue_source).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 +900,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_source) = + 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_source, NUM_JOBS * TASKS_PER_JOB).await?; + assert_no_more_assignments(&*dispatch_queue_source).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 +933,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_source) = + 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_source, NUM_JOBS * TASKS_PER_JOB).await?; + assert_no_more_assignments(&*dispatch_queue_source).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 +984,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_source) = + 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_source, num_assignments).await?; + assert_no_more_assignments(&*dispatch_queue_source).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 +1050,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 +1127,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 +1158,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 +1203,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 +1249,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_source, 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 +1269,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_source + .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..d115e92e3 100644 --- a/components/spider-scheduler/src/dispatch_queue.rs +++ b/components/spider-scheduler/src/dispatch_queue.rs @@ -11,46 +11,9 @@ use tokio::sync::RwLock; use crate::error::SchedulerError; use crate::types::TaskAssignment; -/// The writer side of the dispatching queue used by the scheduler core. -#[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 DispatchQueueSource: Send + Sync { /// Dequeues the next task assignment for an execution manager to execute. /// /// # Parameters @@ -70,8 +33,7 @@ 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. +/// A cloneable writer handle for the dispatching queue, backed by an async channel. /// /// # NOTE /// @@ -83,9 +45,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 +61,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,7 +85,11 @@ impl DispatchQueueSink for DispatchQueueWriter { Ok(()) } - fn size(&self) -> usize { + /// # Returns + /// + /// The current size of the dispatch queue. + #[must_use] + pub fn size(&self) -> usize { self.inner.assignment_sender.len() } } @@ -117,7 +98,8 @@ impl DispatchQueueSink for DispatchQueueWriter { /// an async channel. #[derive(Clone)] pub struct DispatchQueueReader { - inner: Arc, + session_id: Arc>, + assignment_receiver: async_channel::Receiver, } #[async_trait] @@ -125,9 +107,9 @@ impl DispatchQueueSource 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 +117,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 +145,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 +162,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..bc1fdf3a7 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 read 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 (DispatchQueueSource — read handle) //! ▼ //! ┌───────────────────┐ //! │ scheduler service │ ──▶ execution managers (concurrent fan-out) @@ -47,7 +49,6 @@ pub use crate::config::SchedulerConfig; pub use crate::config::ServerConfig; pub use crate::core::SchedulerCore; 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::error::SchedulerError; diff --git a/components/spider-scheduler/src/runtime.rs b/components/spider-scheduler/src/runtime.rs index ec77077a7..59383e2cc 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 read 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 read 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 source. /// * 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_source = core.get_dispatch_queue_source(); + let service = SchedulerServiceState::new(dispatch_queue_source, 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..52937b493 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 source as a trait object, so the runtime can drive it with the read 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; @@ -19,18 +20,12 @@ 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,7 +33,7 @@ impl /// A newly constructed [`SchedulerServiceState`]. #[must_use] pub fn new( - dispatch_source: DispatchQueueSourceType, + dispatch_source: Arc, registry: ExecutionManagerRegistry, scheduler_id: SchedulerId, ) -> Self { @@ -204,12 +199,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_source: Arc, registry: ExecutionManagerRegistry, scheduler_id: SchedulerId, } @@ -255,15 +246,14 @@ 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 [`DispatchQueueSource`] mock backed by a counter. /// /// Each [`DispatchQueueSource::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, + remaining: AtomicUsize, } impl CounterDispatchSource { @@ -271,9 +261,9 @@ mod tests { /// /// A new [`CounterDispatchSource`] that hands out `remaining` assignments before reporting /// an empty queue. - fn new(remaining: usize) -> Self { + const fn new(remaining: usize) -> Self { Self { - remaining: Arc::new(AtomicUsize::new(remaining)), + remaining: AtomicUsize::new(remaining), } } } @@ -370,13 +360,13 @@ mod tests { 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(CounterDispatchSource::new(remaining)), registry, SchedulerId::from(SCHEDULER_ID), ); From acb7bb6f9e6d3073824f1da762452b95ff74eaf8 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Mon, 17 Aug 2026 17:14:02 -0400 Subject: [PATCH 2/4] Minor fix. --- components/spider-scheduler/src/core.rs | 4 ++-- .../src/core_impl/round_robin/implementation.rs | 10 +++++----- .../src/core_impl/round_robin/tests.rs | 5 +++-- components/spider-scheduler/src/dispatch_queue.rs | 4 ++++ components/spider-scheduler/src/lib.rs | 1 + components/spider-scheduler/src/service.rs | 6 +++--- 6 files changed, 18 insertions(+), 12 deletions(-) diff --git a/components/spider-scheduler/src/core.rs b/components/spider-scheduler/src/core.rs index 194195c11..eba5b4afd 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::DispatchQueueSource; +use crate::dispatch_queue::SharedDispatchQueueSource; use crate::error::SchedulerError; use crate::storage_client::SchedulerStorageClient; use crate::types::TaskAssignment; @@ -52,7 +52,7 @@ pub trait SchedulerCore: Send { /// /// A read handle over the core's dispatch structure, for the execution-manager-facing service /// to drain. - fn get_dispatch_queue_source(&self) -> Arc; + fn get_dispatch_queue_source(&self) -> SharedDispatchQueueSource; /// Runs the scheduling loop until `cancellation_token` is triggered. /// 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 ca5453d44..dbcda29cd 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/implementation.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/implementation.rs @@ -26,8 +26,8 @@ use crate::SchedulerStorageClient; use crate::StorageClientError; use crate::TaskAssignment; use crate::core::TaskAssignmentIdIssuer; -use crate::dispatch_queue::DispatchQueueSource; use crate::dispatch_queue::DispatchQueueWriter; +use crate::dispatch_queue::SharedDispatchQueueSource; use crate::dispatch_queue::create_dispatch_queue; /// The configuration of the round-robin scheduler core. @@ -80,7 +80,7 @@ impl RoundRobinConfig { RoundRobinCore { config: self, dispatch_queue_writer, - dispatch_queue_source: Arc::new(dispatch_queue_reader), + dispatch_queue_reader: Arc::new(dispatch_queue_reader), _marker: std::marker::PhantomData, } } @@ -98,7 +98,7 @@ impl RoundRobinConfig { pub struct RoundRobinCore { config: RoundRobinConfig, dispatch_queue_writer: DispatchQueueWriter, - dispatch_queue_source: Arc, + dispatch_queue_reader: SharedDispatchQueueSource, _marker: std::marker::PhantomData, } @@ -108,8 +108,8 @@ impl SchedulerCore { type StorageClient = SchedulerStorageClientType; - fn get_dispatch_queue_source(&self) -> Arc { - Arc::clone(&self.dispatch_queue_source) + fn get_dispatch_queue_source(&self) -> SharedDispatchQueueSource { + Arc::clone(&self.dispatch_queue_reader) } async fn run( 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 8ee558925..2a9c807f9 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/tests.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/tests.rs @@ -33,6 +33,7 @@ use crate::TaskAssignment; use crate::core::TaskAssignmentIdIssuer; use crate::dispatch_queue::DispatchQueueReader; use crate::dispatch_queue::DispatchQueueSource; +use crate::dispatch_queue::SharedDispatchQueueSource; use crate::dispatch_queue::create_dispatch_queue; /// The session used by tests that never bump the session. Must stay equal to the session @@ -309,7 +310,7 @@ fn spawn_scheduler( ) -> ( SchedulerJoinHandle, CancellationToken, - Arc, + SharedDispatchQueueSource, ) { let (handle, cancellation_token, dispatch_queue_source, _reschedule_queue_sender) = spawn_scheduler_with_reschedule(config, storage_client); @@ -334,7 +335,7 @@ fn spawn_scheduler_with_reschedule( ) -> ( SchedulerJoinHandle, CancellationToken, - Arc, + SharedDispatchQueueSource, RescheduleQueueSender, ) { let core = Box::new(config.make_core::()); diff --git a/components/spider-scheduler/src/dispatch_queue.rs b/components/spider-scheduler/src/dispatch_queue.rs index d115e92e3..5abedca70 100644 --- a/components/spider-scheduler/src/dispatch_queue.rs +++ b/components/spider-scheduler/src/dispatch_queue.rs @@ -33,6 +33,10 @@ pub trait DispatchQueueSource: Send + Sync { async fn dequeue(&self, wait_time: Duration) -> Result, SchedulerError>; } +/// The shared read handle over a dispatching queue that a scheduler core hands to the +/// execution-manager-facing service. +pub type SharedDispatchQueueSource = Arc; + /// A cloneable writer handle for the dispatching queue, backed by an async channel. /// /// # NOTE diff --git a/components/spider-scheduler/src/lib.rs b/components/spider-scheduler/src/lib.rs index bc1fdf3a7..b8a736911 100644 --- a/components/spider-scheduler/src/lib.rs +++ b/components/spider-scheduler/src/lib.rs @@ -51,6 +51,7 @@ pub use crate::core::SchedulerCore; pub use crate::dispatch_queue::DispatchQueueReader; pub use crate::dispatch_queue::DispatchQueueSource; pub use crate::dispatch_queue::DispatchQueueWriter; +pub use crate::dispatch_queue::SharedDispatchQueueSource; pub use crate::error::SchedulerError; pub use crate::error::SchedulerRuntimeError; pub use crate::error::SchedulerServiceError; diff --git a/components/spider-scheduler/src/service.rs b/components/spider-scheduler/src/service.rs index 52937b493..ec1983881 100644 --- a/components/spider-scheduler/src/service.rs +++ b/components/spider-scheduler/src/service.rs @@ -14,7 +14,7 @@ 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::SharedDispatchQueueSource; use crate::error::SchedulerServiceError; use crate::execution_manager_registry::ExecutionManagerRegistry; use crate::types::TaskAssignment; @@ -33,7 +33,7 @@ impl SchedulerServiceState { /// A newly constructed [`SchedulerServiceState`]. #[must_use] pub fn new( - dispatch_source: Arc, + dispatch_source: SharedDispatchQueueSource, registry: ExecutionManagerRegistry, scheduler_id: SchedulerId, ) -> Self { @@ -200,7 +200,7 @@ impl SchedulerServiceState { /// The shared inner state of [`SchedulerServiceState`]. struct SchedulerServiceStateInner { - dispatch_source: Arc, + dispatch_source: SharedDispatchQueueSource, registry: ExecutionManagerRegistry, scheduler_id: SchedulerId, } From 61b17e54452f154a115ee6bdfa64dcad77e1ea9b Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Mon, 17 Aug 2026 17:23:42 -0400 Subject: [PATCH 3/4] refactor(spider-scheduler): Move dispatch queue ownership into the scheduler core: * Make the core own and construct its dispatch queue instead of receiving a pre-built sink as a parameter to `run`, so that a core is free to own a dispatch structure that is not a single channel. `SchedulerCore` loses its `Sink` associated type and gains `get_dispatch_queue_handle`, and the runtime now builds the core first and takes the handle from it. * Delete the `DispatchQueueSink` trait, whose sole implementor became a concrete type once the core owned its writer, moving `enqueue`, `bump_session_id` and `size` to an inherent `impl DispatchQueueWriter`. * Rename `DispatchQueueSource` to `DispatchQueueHandle`, since with the sink gone it named one half of a pair that no longer exists, and add a `SharedDispatchQueueHandle` alias for its `Arc`. Drop `Clone` from the trait's supertraits so it is object-safe, and flatten `DispatchQueueReader` so the handle carries a single level of indirection. * Drop the now-unused `SchedulerConfig::dispatch_queue_capacity`, and the `DispatchQueueSourceType` type parameter from `SchedulerServiceState` and `GrpcSchedulerService`. --- components/spider-scheduler/src/core.rs | 8 +-- .../core_impl/round_robin/implementation.rs | 8 +-- .../src/core_impl/round_robin/tests.rs | 64 +++++++++---------- .../spider-scheduler/src/dispatch_queue.rs | 13 ++-- components/spider-scheduler/src/lib.rs | 8 +-- components/spider-scheduler/src/runtime.rs | 14 ++-- components/spider-scheduler/src/service.rs | 42 ++++++------ 7 files changed, 80 insertions(+), 77 deletions(-) diff --git a/components/spider-scheduler/src/core.rs b/components/spider-scheduler/src/core.rs index eba5b4afd..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::SharedDispatchQueueSource; +use crate::dispatch_queue::SharedDispatchQueueHandle; use crate::error::SchedulerError; use crate::storage_client::SchedulerStorageClient; use crate::types::TaskAssignment; @@ -50,9 +50,9 @@ pub trait SchedulerCore: Send { /// # Returns /// - /// A read handle over the core's dispatch structure, for the execution-manager-facing service - /// to drain. - fn get_dispatch_queue_source(&self) -> SharedDispatchQueueSource; + /// 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. /// 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 dbcda29cd..6d5a362ad 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/implementation.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/implementation.rs @@ -27,7 +27,7 @@ use crate::StorageClientError; use crate::TaskAssignment; use crate::core::TaskAssignmentIdIssuer; use crate::dispatch_queue::DispatchQueueWriter; -use crate::dispatch_queue::SharedDispatchQueueSource; +use crate::dispatch_queue::SharedDispatchQueueHandle; use crate::dispatch_queue::create_dispatch_queue; /// The configuration of the round-robin scheduler core. @@ -98,7 +98,7 @@ impl RoundRobinConfig { pub struct RoundRobinCore { config: RoundRobinConfig, dispatch_queue_writer: DispatchQueueWriter, - dispatch_queue_reader: SharedDispatchQueueSource, + dispatch_queue_reader: SharedDispatchQueueHandle, _marker: std::marker::PhantomData, } @@ -108,8 +108,8 @@ impl SchedulerCore { type StorageClient = SchedulerStorageClientType; - fn get_dispatch_queue_source(&self) -> SharedDispatchQueueSource { - Arc::clone(&self.dispatch_queue_reader) + fn get_dispatch_queue_handle(&self) -> SharedDispatchQueueHandle { + self.dispatch_queue_reader.clone() } async fn run( 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 2a9c807f9..9f62dee6b 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/tests.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/tests.rs @@ -31,9 +31,9 @@ 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::DispatchQueueSource; -use crate::dispatch_queue::SharedDispatchQueueSource; +use crate::dispatch_queue::SharedDispatchQueueHandle; use crate::dispatch_queue::create_dispatch_queue; /// The session used by tests that never bump the session. Must stay equal to the session @@ -303,18 +303,18 @@ fn make_reschedule_assignment( /// /// * The join handle yielding the scheduler's exit result. /// * The cancellation token that stops the scheduler. -/// * The read handle over the dispatch queue the core owns, playing the execution manager's role. +/// * The handle over the dispatch queue the core owns, playing the execution manager's role. fn spawn_scheduler( config: RoundRobinConfig, storage_client: MockStorageClient, ) -> ( SchedulerJoinHandle, CancellationToken, - SharedDispatchQueueSource, + SharedDispatchQueueHandle, ) { - let (handle, cancellation_token, dispatch_queue_source, _reschedule_queue_sender) = + let (handle, cancellation_token, dispatch_queue_handle, _reschedule_queue_sender) = spawn_scheduler_with_reschedule(config, storage_client); - (handle, cancellation_token, dispatch_queue_source) + (handle, cancellation_token, dispatch_queue_handle) } /// Spawns the scheduler's public run loop as a background task, exposing the reschedule-queue @@ -326,7 +326,7 @@ fn spawn_scheduler( /// /// * The join handle yielding the scheduler's exit result. /// * The cancellation token that stops the scheduler. -/// * The read handle over the dispatch queue the core owns, playing the execution manager's role. +/// * 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( @@ -335,11 +335,11 @@ fn spawn_scheduler_with_reschedule( ) -> ( SchedulerJoinHandle, CancellationToken, - SharedDispatchQueueSource, + SharedDispatchQueueHandle, RescheduleQueueSender, ) { let core = Box::new(config.make_core::()); - let dispatch_queue_source = core.get_dispatch_queue_source(); + 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(); @@ -355,7 +355,7 @@ fn spawn_scheduler_with_reschedule( ( handle, cancellation_token, - dispatch_queue_source, + dispatch_queue_handle, reschedule_queue_sender, ) } @@ -372,9 +372,9 @@ 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. +/// * Forwards [`DispatchQueueHandle::dequeue`]'s return values on failure. async fn drain_n( - dispatch_queue_source: &dyn DispatchQueueSource, + dispatch_queue_handle: &dyn DispatchQueueHandle, n: usize, ) -> anyhow::Result> { const DEQUEUE_WAIT: Duration = Duration::from_millis(100); @@ -387,7 +387,7 @@ async fn drain_n( assignments.len(), ); } - if let Some(assignment) = dispatch_queue_source.dequeue(DEQUEUE_WAIT).await? { + if let Some(assignment) = dispatch_queue_handle.dequeue(DEQUEUE_WAIT).await? { assignments.push(assignment); } } @@ -401,16 +401,16 @@ async fn drain_n( /// /// Returns an error if: /// -/// * Forwards [`DispatchQueueSource::dequeue`]'s return values on failure. +/// * Forwards [`DispatchQueueHandle::dequeue`]'s return values on failure. /// /// # Panics /// /// Panics if an assignment arrives within the observation window. async fn assert_no_more_assignments( - dispatch_queue_source: &dyn DispatchQueueSource, + dispatch_queue_handle: &dyn DispatchQueueHandle, ) -> anyhow::Result<()> { const OBSERVATION_WINDOW: Duration = Duration::from_secs(1); - let unexpected_assignment = dispatch_queue_source.dequeue(OBSERVATION_WINDOW).await?; + let unexpected_assignment = dispatch_queue_handle.dequeue(OBSERVATION_WINDOW).await?; assert_eq!(unexpected_assignment, None); Ok(()) } @@ -588,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, @@ -619,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 /// @@ -869,11 +869,11 @@ async fn single_capacity_pool_schedules_jobs_serially() -> anyhow::Result<()> { ); let config = make_config(1, DISPATCH_QUEUE_CAPACITY); - let (scheduler_handle, cancellation_token, dispatch_queue_source) = + let (scheduler_handle, cancellation_token, dispatch_queue_handle) = spawn_scheduler(config, storage_client); - let assignments = drain_n(&*dispatch_queue_source, NUM_JOBS * TASKS_PER_JOB).await?; - assert_no_more_assignments(&*dispatch_queue_source).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 @@ -902,11 +902,11 @@ async fn active_jobs_dispatch_in_round_robin_order() -> anyhow::Result<()> { ); let config = make_config(NUM_JOBS, DISPATCH_QUEUE_CAPACITY); - let (scheduler_handle, cancellation_token, dispatch_queue_source) = + let (scheduler_handle, cancellation_token, dispatch_queue_handle) = spawn_scheduler(config, storage_client); - let assignments = drain_n(&*dispatch_queue_source, NUM_JOBS * TASKS_PER_JOB).await?; - assert_no_more_assignments(&*dispatch_queue_source).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 @@ -935,11 +935,11 @@ async fn pending_jobs_promote_and_schedule_round_robin() -> anyhow::Result<()> { ); let config = make_config(ACTIVE_JOB_QUEUE_CAPACITY, DISPATCH_QUEUE_CAPACITY); - let (scheduler_handle, cancellation_token, dispatch_queue_source) = + let (scheduler_handle, cancellation_token, dispatch_queue_handle) = spawn_scheduler(config, storage_client); - let assignments = drain_n(&*dispatch_queue_source, NUM_JOBS * TASKS_PER_JOB).await?; - assert_no_more_assignments(&*dispatch_queue_source).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); @@ -986,13 +986,13 @@ async fn commit_drains_each_cycle_cleanup_dispatches_once() -> anyhow::Result<() ); let config = make_config(NUM_ACTIVE_JOBS, DISPATCH_QUEUE_CAPACITY); - let (scheduler_handle, cancellation_token, dispatch_queue_source) = + 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(&*dispatch_queue_source, num_assignments).await?; - assert_no_more_assignments(&*dispatch_queue_source).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 @@ -1250,7 +1250,7 @@ async fn randomly_rescheduled_assignments_are_eventually_redispatched() -> anyho make_ready_batch(&jobs, TASKS_PER_JOB, 0), ); - let (scheduler_handle, cancellation_token, dispatch_queue_source, reschedule_sender) = + let (scheduler_handle, cancellation_token, dispatch_queue_handle, reschedule_sender) = spawn_scheduler_with_reschedule( make_config(NUM_JOBS, DISPATCH_QUEUE_CAPACITY), storage_client, @@ -1270,7 +1270,7 @@ async fn randomly_rescheduled_assignments_are_eventually_redispatched() -> anyho completed.len(), ); } - let Some(assignment) = dispatch_queue_source + let Some(assignment) = dispatch_queue_handle .dequeue(Duration::from_millis(100)) .await? else { diff --git a/components/spider-scheduler/src/dispatch_queue.rs b/components/spider-scheduler/src/dispatch_queue.rs index 5abedca70..a5b020d91 100644 --- a/components/spider-scheduler/src/dispatch_queue.rs +++ b/components/spider-scheduler/src/dispatch_queue.rs @@ -11,9 +11,10 @@ use tokio::sync::RwLock; use crate::error::SchedulerError; use crate::types::TaskAssignment; -/// The reader side of the dispatching queue, drained by the execution-manager-facing service. +/// 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 DispatchQueueSource: Send + Sync { +pub trait DispatchQueueHandle: Send + Sync { /// Dequeues the next task assignment for an execution manager to execute. /// /// # Parameters @@ -33,9 +34,9 @@ pub trait DispatchQueueSource: Send + Sync { async fn dequeue(&self, wait_time: Duration) -> Result, SchedulerError>; } -/// The shared read handle over a dispatching queue that a scheduler core hands to the +/// The shared handle over a dispatching queue that a scheduler core hands to the /// execution-manager-facing service. -pub type SharedDispatchQueueSource = Arc; +pub type SharedDispatchQueueHandle = Arc; /// A cloneable writer handle for the dispatching queue, backed by an async channel. /// @@ -98,7 +99,7 @@ impl DispatchQueueWriter { } } -/// 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 { @@ -107,7 +108,7 @@ pub struct DispatchQueueReader { } #[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. diff --git a/components/spider-scheduler/src/lib.rs b/components/spider-scheduler/src/lib.rs index b8a736911..7d692e972 100644 --- a/components/spider-scheduler/src/lib.rs +++ b/components/spider-scheduler/src/lib.rs @@ -8,7 +8,7 @@ //! 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. The core owns and constructs its own dispatching queue, -//! and exposes only a read handle to the service: +//! and exposes only a handle to the service: //! //! ```text //! storage ── authoritative inbound queue (owned by the storage layer, not this crate) @@ -26,7 +26,7 @@ //! │ └───────────────────────────────────────────┘ │ //! └───────────────────────────────────────────────────┘ //! │ -//! │ dequeue (DispatchQueueSource — read handle) +//! │ dequeue (DispatchQueueHandle) //! ▼ //! ┌───────────────────┐ //! │ scheduler service │ ──▶ execution managers (concurrent fan-out) @@ -48,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::DispatchQueueSource; pub use crate::dispatch_queue::DispatchQueueWriter; -pub use crate::dispatch_queue::SharedDispatchQueueSource; +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 59383e2cc..9d6e2e286 100644 --- a/components/spider-scheduler/src/runtime.rs +++ b/components/spider-scheduler/src/runtime.rs @@ -1,8 +1,8 @@ //! The scheduler runtime. //! //! This module registers the scheduler with the storage service, builds the scheduler core and -//! wires the execution-manager-facing service to the read handle of the dispatch queue the core -//! owns, hands the core the reschedule queue reader, and spawns the core's scheduling loop as a +//! 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. @@ -85,8 +85,8 @@ impl Runtime { /// Creates a scheduler runtime from the given configuration and storage client. /// /// Registers this scheduler with the storage service, builds the scheduler core and wires the -/// execution-manager-facing service to the read 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 +/// 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 @@ -98,7 +98,7 @@ impl Runtime { /// A tuple on success, containing: /// /// * The newly created runtime instance. -/// * The execution-manager-facing scheduler service, built over the core's dispatch queue source. +/// * 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 @@ -132,8 +132,8 @@ pub async fn create_runtime(); - let dispatch_queue_source = core.get_dispatch_queue_source(); - let service = SchedulerServiceState::new(dispatch_queue_source, registry, scheduler_id); + 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 { diff --git a/components/spider-scheduler/src/service.rs b/components/spider-scheduler/src/service.rs index ec1983881..582e61950 100644 --- a/components/spider-scheduler/src/service.rs +++ b/components/spider-scheduler/src/service.rs @@ -3,8 +3,8 @@ //! [`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 -//! holds its dispatch source as a trait object, so the runtime can drive it with the read handle of -//! whichever dispatch structure the scheduler core owns in production, or with a mock in tests, +//! 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; @@ -14,7 +14,7 @@ use spider_core::types::id::ExecutionManagerId; use spider_core::types::id::SchedulerId; use spider_core::types::scheduler::TaskAssignmentRecord; -use crate::dispatch_queue::SharedDispatchQueueSource; +use crate::dispatch_queue::SharedDispatchQueueHandle; use crate::error::SchedulerServiceError; use crate::execution_manager_registry::ExecutionManagerRegistry; use crate::types::TaskAssignment; @@ -33,13 +33,13 @@ impl SchedulerServiceState { /// A newly constructed [`SchedulerServiceState`]. #[must_use] pub fn new( - dispatch_source: SharedDispatchQueueSource, + dispatch_queue_handle: SharedDispatchQueueHandle, registry: ExecutionManagerRegistry, scheduler_id: SchedulerId, ) -> Self { Self { inner: Arc::new(SchedulerServiceStateInner { - dispatch_source, + dispatch_queue_handle, registry, scheduler_id, }), @@ -73,7 +73,9 @@ impl SchedulerServiceState { /// /// 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, @@ -90,7 +92,7 @@ impl SchedulerServiceState { 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(), @@ -200,7 +202,7 @@ impl SchedulerServiceState { /// The shared inner state of [`SchedulerServiceState`]. struct SchedulerServiceStateInner { - dispatch_source: SharedDispatchQueueSource, + dispatch_queue_handle: SharedDispatchQueueHandle, registry: ExecutionManagerRegistry, scheduler_id: SchedulerId, } @@ -228,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. @@ -246,21 +248,21 @@ 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 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. - struct CounterDispatchSource { + struct CounterDispatchQueueHandle { remaining: AtomicUsize, } - impl CounterDispatchSource { + impl CounterDispatchQueueHandle { /// # Returns /// - /// A new [`CounterDispatchSource`] that hands out `remaining` assignments before reporting - /// an empty queue. + /// A new [`CounterDispatchQueueHandle`] that hands out `remaining` assignments before + /// reporting an empty queue. const fn new(remaining: usize) -> Self { Self { remaining: AtomicUsize::new(remaining), @@ -269,7 +271,7 @@ mod tests { } #[async_trait] - impl DispatchQueueSource for CounterDispatchSource { + impl DispatchQueueHandle for CounterDispatchQueueHandle { async fn dequeue( &self, _wait_time: Duration, @@ -353,8 +355,8 @@ 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( @@ -366,7 +368,7 @@ mod tests { ) { let (registry, reschedule_queue_receiver, cancellation_token) = build_registry(); let service = SchedulerServiceState::new( - Arc::new(CounterDispatchSource::new(remaining)), + Arc::new(CounterDispatchQueueHandle::new(remaining)), registry, SchedulerId::from(SCHEDULER_ID), ); From 5d55d6e04b4cddfc94000008fec08125d4888ab5 Mon Sep 17 00:00:00 2001 From: LinZhihao-723 Date: Tue, 18 Aug 2026 13:09:10 -0400 Subject: [PATCH 4/4] Remove docstring. --- components/spider-scheduler/src/dispatch_queue.rs | 3 --- 1 file changed, 3 deletions(-) diff --git a/components/spider-scheduler/src/dispatch_queue.rs b/components/spider-scheduler/src/dispatch_queue.rs index a5b020d91..fc300bc6d 100644 --- a/components/spider-scheduler/src/dispatch_queue.rs +++ b/components/spider-scheduler/src/dispatch_queue.rs @@ -90,9 +90,6 @@ impl DispatchQueueWriter { Ok(()) } - /// # Returns - /// - /// The current size of the dispatch queue. #[must_use] pub fn size(&self) -> usize { self.inner.assignment_sender.len()