diff --git a/Cargo.lock b/Cargo.lock index fb506b56d..9753bb931 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2026,6 +2026,7 @@ dependencies = [ "async-trait", "clap", "dashmap", + "rand 0.9.4", "serde", "spider-core", "spider-proto-rust", diff --git a/components/spider-core/src/types/scheduler.rs b/components/spider-core/src/types/scheduler.rs index a86f3a403..f17914469 100644 --- a/components/spider-core/src/types/scheduler.rs +++ b/components/spider-core/src/types/scheduler.rs @@ -29,6 +29,9 @@ pub struct TaskAssignment { /// The task to dispatch. pub task_id: TaskId, + + /// The scheduler's view of storage's session when the assignment was produced. + pub session_id: SessionId, } /// A record of a task assignment previously produced by the scheduler. @@ -50,8 +53,7 @@ impl TaskAssignmentRecord { /// A task assignment handed to the execution manager by the scheduler. /// -/// `session_id` is the scheduler's view of storage's session at the moment the assignment was -/// produced. The execution manager pins this exact value on every subsequent storage call for the +/// The execution manager pins `task_assignment.session_id` on every subsequent storage call for the /// attempt. #[derive(Debug, Clone, Copy, PartialEq, Eq)] pub struct SchedulerResponse { @@ -60,7 +62,4 @@ pub struct SchedulerResponse { /// The scheduler that produced the assignment. pub scheduler_id: SchedulerId, - - /// The scheduler's view of storage's session when the assignment was produced. - pub session_id: SessionId, } diff --git a/components/spider-execution-manager/src/runtime.rs b/components/spider-execution-manager/src/runtime.rs index b3bad11a2..33ba75f26 100644 --- a/components/spider-execution-manager/src/runtime.rs +++ b/components/spider-execution-manager/src/runtime.rs @@ -302,16 +302,16 @@ impl< }; tracing::info!( - bundle_session = response.session_id, + bundle_session = response.task_assignment.session_id, job_id = ? response.task_assignment.job_id, task_id = ? response.task_assignment.task_id, "Received a new task assignment from the scheduler." ); let current_session = self.session_tracker.current(); - if response.session_id < current_session { + if response.task_assignment.session_id < current_session { tracing::warn!( - bundle_session = response.session_id, + bundle_session = response.task_assignment.session_id, current_session, job_id = ? response.task_assignment.job_id, task_id = ? response.task_assignment.task_id, @@ -320,9 +320,9 @@ impl< self.mark_consume(&response); continue; } - if response.session_id > current_session { + if response.task_assignment.session_id > current_session { tracing::info!( - new_session = response.session_id, + new_session = response.task_assignment.session_id, "Observed a newer session via the scheduler. Refreshing liveness." ); self.liveness_handle.refresh().await; @@ -355,9 +355,9 @@ impl< })?; let current_session = self.session_tracker.current(); - if response.session_id < current_session { + if response.task_assignment.session_id < current_session { tracing::warn!( - bundle_session = response.session_id, + bundle_session = response.task_assignment.session_id, current_session, job_id = ? response.task_assignment.job_id, task_id = ? response.task_assignment.task_id, @@ -375,7 +375,7 @@ impl< job: response.task_assignment.job_id, task: response.task_assignment.task_id, task_instance_id, - session: response.session_id, + session: response.task_assignment.session_id, }, outcome, )); @@ -413,7 +413,7 @@ impl< response.task_assignment.job_id, response.task_assignment.task_id, self.em_id, - response.session_id, + response.task_assignment.session_id, ) => result, }; @@ -425,7 +425,7 @@ impl< Err(err) => match &err { StorageResponseError::StaleSession(message) => { tracing::warn!( - bundle_session = response.session_id, + bundle_session = response.task_assignment.session_id, error = % message, job_id = ? response.task_assignment.job_id, task_id = ? response.task_assignment.task_id, diff --git a/components/spider-proto-rust/src/assignment.rs b/components/spider-proto-rust/src/assignment.rs index 88dded6f7..29a38b32c 100644 --- a/components/spider-proto-rust/src/assignment.rs +++ b/components/spider-proto-rust/src/assignment.rs @@ -46,9 +46,9 @@ impl TryFrom for Option { resource_group_id: ResourceGroupId::from(assignment.resource_group_id), job_id: JobId::from(assignment.job_id), task_id, + session_id: assignment.session_id, }, scheduler_id: SchedulerId::from(assignment.scheduler_id), - session_id: assignment.session_id, })) } Some(next_task_response::Result::NoTask(common::Void {})) => Ok(None), @@ -117,7 +117,7 @@ mod tests { assert_eq!(scheduler_response.task_assignment.job_id, JobId::from(3)); assert_eq!(scheduler_response.task_assignment.task_id, TaskId::Index(7)); assert_eq!(scheduler_response.scheduler_id, SchedulerId::from(4)); - assert_eq!(scheduler_response.session_id, 5); + assert_eq!(scheduler_response.task_assignment.session_id, 5); } #[test] diff --git a/components/spider-scheduler/Cargo.toml b/components/spider-scheduler/Cargo.toml index 26b8a7ef3..256a269b9 100644 --- a/components/spider-scheduler/Cargo.toml +++ b/components/spider-scheduler/Cargo.toml @@ -31,5 +31,6 @@ tracing = { version = "0.1.41", default-features = false, features = ["std"] } [dev-dependencies] anyhow = "1.0.102" dashmap = "6.1.0" +rand = "0.9.1" tokio = { version = "1.52.3", features = ["macros", "rt-multi-thread"] } tokio-util = { version = "0.7.18", features = ["rt"] } diff --git a/components/spider-scheduler/src/core.rs b/components/spider-scheduler/src/core.rs index 8dee0dc07..4a2bda858 100644 --- a/components/spider-scheduler/src/core.rs +++ b/components/spider-scheduler/src/core.rs @@ -9,6 +9,7 @@ use spider_core::types::id::TaskAssignmentId; use crate::dispatch_queue::DispatchQueueSink; use crate::error::SchedulerError; use crate::storage_client::SchedulerStorageClient; +use crate::types::TaskAssignment; /// Single-source ID issuer for creating globally unique IDs for task assignments. pub struct TaskAssignmentIdIssuer { @@ -63,6 +64,8 @@ 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 /// assignments. /// * `cancellation_token` - The token to signal the scheduling loop to stop. @@ -74,6 +77,7 @@ pub trait SchedulerCore: Send { 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, ) -> Result<(), SchedulerError>; 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 8c02406f7..dd9a573a7 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/implementation.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/implementation.rs @@ -114,6 +114,7 @@ impl< self: Box, storage_client: Self::StorageClient, sink: Self::Sink, + reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver, id_issuer: TaskAssignmentIdIssuer, cancellation_token: CancellationToken, ) -> Result<(), SchedulerError> { @@ -121,6 +122,7 @@ impl< SessionId::default(), storage_client, sink, + reschedule_queue_reader, id_issuer, cancellation_token, self.config, @@ -202,6 +204,7 @@ pub(super) struct RoundRobin< pub(super) finalizing_job_queue: VecDeque<(JobId, Instant)>, pub(super) inbound_queue_reader: AsyncInboundQueueReader, + pub(super) reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver, } impl< @@ -220,6 +223,7 @@ impl< storage_session_id: SessionId, storage_client: SchedulerStorageClientType, sink: DispatchQueueSinkType, + reschedule_queue_reader: tokio::sync::mpsc::UnboundedReceiver, id_issuer: TaskAssignmentIdIssuer, cancellation_token: CancellationToken, config: RoundRobinConfig, @@ -254,21 +258,25 @@ impl< finalizing_jobs, finalizing_job_queue, inbound_queue_reader, + reschedule_queue_reader, } } - /// Executes a single scheduling tick: consumes any completed inbound poll, then makes - /// scheduling decisions to fill the dispatch queue. + /// Executes a single scheduling tick: consumes any completed inbound poll, re-injects + /// assignments returned by lost execution managers, then makes scheduling decisions to fill the + /// dispatch queue. /// /// # Errors /// /// Returns an error if: /// /// * Forwards [`Self::consume_inbound_poll_result`]'s return values on failure. + /// * Forwards [`Self::reschedule`]'s return values on failure. /// * Forwards [`Self::make_schedule_decisions`]'s return values on failure. pub(super) async fn tick(&mut self) -> Result<(), SchedulerError> { tracing::info!("Starting scheduling tick."); self.consume_inbound_poll_result().await?; + self.reschedule()?; self.make_schedule_decisions().await?; self.retire_expired_finalizing_jobs(); Ok(()) @@ -488,6 +496,78 @@ impl< Ok(()) } + /// Re-injects assignments recovered from lost execution managers into the scheduler. + /// + /// The assignments are grouped by task kind and processed through the same + /// commit-ready, cleanup-ready, and ready paths as inbound entries. This ensures that + /// finalization semantics and buffered-task deduplication are applied consistently. + /// + /// Rescheduled assignments are enqueued without respecting the internal buffer's capacity + /// limit. As a result, the buffer may temporarily contain more assignments than its configured + /// capacity. + /// + /// # Errors + /// + /// Returns an error if: + /// + /// * [`SchedulerError::Internal`] if the re-schedule queue reader fails to receive a buffered + /// assignment. + /// * Forwards [`Self::enqueue_commit_ready_entries`]'s return values on failure. + /// * Forwards [`Self::enqueue_cleanup_ready_entries`]'s return values on failure. + fn reschedule(&mut self) -> Result<(), SchedulerError> { + if self.reschedule_queue_reader.is_empty() { + return Ok(()); + } + + // Snapshot the size so assignments pushed concurrently during the drain are deferred to the + // next tick instead of extending this loop unboundedly. + let mut reschedule_queue_size = self.reschedule_queue_reader.len(); + let mut ready_entries = Vec::new(); + let mut commit_ready_entries = Vec::new(); + let mut cleanup_ready_entries = Vec::new(); + while reschedule_queue_size > 0 { + let assignment = self.reschedule_queue_reader.try_recv().map_err(|e| { + tracing::error!( + err = % e, + "Reschedule queue reader failed to receive a message." + ); + SchedulerError::Internal( + "reschedule queue reader failed to receive a message".to_string(), + ) + })?; + reschedule_queue_size -= 1; + + if assignment.session_id < self.storage_session_id { + continue; + } + + let entry = InboundEntry { + resource_group_id: assignment.resource_group_id, + job_id: assignment.job_id, + task_id: assignment.task_id, + }; + match &assignment.task_id { + TaskId::Index(_) => { + ready_entries.push(entry); + } + TaskId::Commit => { + commit_ready_entries.push(entry); + } + TaskId::Cleanup => { + cleanup_ready_entries.push(entry); + } + } + } + + // Load commit-ready tasks and cleanup-ready tasks first to avoid loading a job that is + // already finalizing. + self.enqueue_commit_ready_entries(commit_ready_entries)?; + self.enqueue_cleanup_ready_entries(cleanup_ready_entries)?; + self.enqueue_ready_entries(ready_entries); + + Ok(()) + } + /// Enqueues polled commit-ready entries: each entry's job is marked finalizing, queued for a /// commit-task assignment, and removed from the active or pending set. /// @@ -732,6 +812,7 @@ impl< job_id, resource_group_id, task_id: TaskId::Cleanup, + session_id: self.storage_session_id, }) .await?; self.buffered_tasks.remove(&(job_id, TaskId::Cleanup)); @@ -752,6 +833,7 @@ impl< job_id, resource_group_id, task_id: TaskId::Commit, + session_id: self.storage_session_id, }) .await?; self.buffered_tasks.remove(&(job_id, TaskId::Commit)); @@ -771,6 +853,7 @@ impl< job_id, resource_group_id: job_entry.resource_group_id, task_id, + session_id: self.storage_session_id, }) .await?; self.buffered_tasks.remove(&(job_id, task_id)); @@ -800,6 +883,8 @@ impl< /// /// * Forwards [`AsyncInboundQueueReader::start`]'s return values on failure. fn start_inbound_poll(&mut self) -> Result<(), SchedulerError> { + // The reschedule path can load the ready buffers beyond their configured capacity, so + // `saturating_sub` floors the remaining poll budget at zero instead of underflowing. let num_commit_ready_tasks = self.commit_ready_jobs.len(); let num_cleanup_ready_tasks = self.cleanup_ready_jobs.len(); let max_commit_ready_entries = self 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 8a9d8fc6a..3364fe557 100644 --- a/components/spider-scheduler/src/core_impl/round_robin/tests.rs +++ b/components/spider-scheduler/src/core_impl/round_robin/tests.rs @@ -18,6 +18,7 @@ use spider_core::types::id::JobId; use spider_core::types::id::ResourceGroupId; use spider_core::types::id::SchedulerId; use spider_core::types::id::SessionId; +use spider_core::types::id::TaskAssignmentId; use spider_core::types::id::TaskId; use tokio_util::sync::CancellationToken; @@ -263,6 +264,29 @@ fn make_finalizing_batch(jobs: &[(JobId, ResourceGroupId)], task_id: TaskId) -> .collect() } +/// Builds a [`TaskAssignment`] as the execution-manager registry would push onto the re-schedule +/// queue for a lost assignment. The `id` is drawn from the top of the id space so it can never +/// collide with the sequential ids the scheduler's own issuer hands out, letting a test prove an +/// assignment was re-issued. +/// +/// # Returns +/// +/// The reschedule-queue [`TaskAssignment`] for the given job, task, and session. +fn make_reschedule_assignment( + job: (JobId, ResourceGroupId), + task_id: TaskId, + session_id: SessionId, + distinct_id: u64, +) -> TaskAssignment { + TaskAssignment { + id: TaskAssignmentId::from(u64::MAX - distinct_id), + resource_group_id: job.1, + job_id: job.0, + task_id, + session_id, + } +} + /// Spawns the scheduler's public run loop as a background task. /// /// # Returns @@ -278,20 +302,47 @@ fn spawn_scheduler( ) -> ( tokio::task::JoinHandle>, CancellationToken, +) { + let (handle, cancellation_token, _reschedule_queue_sender) = + spawn_scheduler_with_reschedule(config, storage_client, sink); + (handle, cancellation_token) +} + +/// Spawns the scheduler's public run loop as a background task, exposing the reschedule-queue +/// sender. +/// +/// # Returns +/// +/// A tuple containing: +/// +/// * 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. +fn spawn_scheduler_with_reschedule( + config: RoundRobinConfig, + storage_client: MockStorageClient, + sink: DispatchQueueWriter, +) -> ( + tokio::task::JoinHandle>, + CancellationToken, + tokio::sync::mpsc::UnboundedSender, ) { let core = Box::new(config.make_core()); 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) + (handle, cancellation_token, reschedule_queue_sender) } /// Drains exactly `n` task assignments from the dispatch queue, playing the worker pool's role. @@ -317,7 +368,7 @@ async fn drain_n(reader: &DispatchQueueReader, n: usize) -> anyhow::Result TestScheduler { - RoundRobin::new( + make_scheduler_with_reschedule(config, storage_client, sink).0 +} + +/// # 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. +fn make_scheduler_with_reschedule( + config: RoundRobinConfig, + storage_client: MockStorageClient, + sink: DispatchQueueWriter, +) -> ( + TestScheduler, + tokio::sync::mpsc::UnboundedSender, +) { + let (reschedule_queue_sender, reschedule_queue_reader) = tokio::sync::mpsc::unbounded_channel(); + let scheduler = RoundRobin::new( DEFAULT_SESSION_ID, storage_client, sink, + reschedule_queue_reader, TaskAssignmentIdIssuer::new(), CancellationToken::new(), config, - ) + ); + (scheduler, reschedule_queue_sender) } /// Ticks the scheduler until `predicate` holds on its state. @@ -496,7 +566,7 @@ async fn tick_and_drain_n( scheduler: &mut TestScheduler, reader: &DispatchQueueReader, n: usize, -) -> anyhow::Result> { +) -> anyhow::Result> { let deadline = tokio::time::Instant::now() + DRAIN_DEADLINE; let mut assignments = Vec::with_capacity(n); while assignments.len() < n { @@ -507,8 +577,8 @@ async fn tick_and_drain_n( ); } scheduler.tick().await?; - while let Some((session_id, assignment)) = reader.dequeue(Duration::ZERO).await? { - assignments.push((session_id, assignment)); + while let Some(assignment) = reader.dequeue(Duration::ZERO).await? { + assignments.push(assignment); } tokio::task::yield_now().await; } @@ -669,11 +739,7 @@ async fn assert_finalizing_ready_drops_jobs(finalizing_task_id: TaskId) -> anyho let num_assignments = NUM_PRE_FREEZE_ASSIGNMENTS + NUM_FINALIZED_JOBS + (TASKS_PER_JOB - 1) + TASKS_PER_JOB; let assignments: Vec = - tick_and_drain_n(&mut scheduler, &reader, num_assignments) - .await? - .into_iter() - .map(|(_session_id, assignment)| assignment) - .collect(); + tick_and_drain_n(&mut scheduler, &reader, num_assignments).await?; assert_no_further_assignments(&mut scheduler, &reader).await?; assert_eq!(scheduler.buffered_tasks.len(), 0); @@ -756,11 +822,7 @@ async fn assert_finalizing_ready_drops_jobs(finalizing_task_id: TaskId) -> anyho late_batch.extend(make_ready_batch(&canary_jobs, TASKS_PER_JOB, 0)); storage_client.push_ready_batch(DEFAULT_SESSION_ID, late_batch); - let late_assignments: Vec<_> = tick_and_drain_n(&mut scheduler, &reader, TASKS_PER_JOB) - .await? - .into_iter() - .map(|(_session_id, assignment)| assignment) - .collect(); + let late_assignments = tick_and_drain_n(&mut scheduler, &reader, TASKS_PER_JOB).await?; assert_strict_rotation(&late_assignments, &canary_jobs, TASKS_PER_JOB); assert_no_further_assignments(&mut scheduler, &reader).await?; @@ -1022,18 +1084,220 @@ async fn session_bump_clears_buffered_tasks() -> anyhow::Result<()> { // draining yields exactly the new jobs' tasks in strict rotation, each paired with the new // session. let num_new_assignments = new_jobs.len() * NEW_TASKS_PER_JOB; - let session_stamped = tick_and_drain_n(&mut scheduler, &reader, num_new_assignments).await?; + let assignments = tick_and_drain_n(&mut scheduler, &reader, num_new_assignments).await?; assert_no_further_assignments(&mut scheduler, &reader).await?; - for &(session_id, _) in &session_stamped { - assert_eq!(session_id, NEW_SESSION_ID); + for assignment in &assignments { + assert_eq!(assignment.session_id, NEW_SESSION_ID); } - let assignments: Vec = session_stamped - .into_iter() - .map(|(_session_id, assignment)| assignment) - .collect(); assert_strict_rotation(&assignments, &new_jobs, NEW_TASKS_PER_JOB); Ok(()) } + +#[tokio::test] +async fn reschedule_ready_task_redispatches_with_current_session() -> anyhow::Result<()> { + const ACTIVE_JOB_QUEUE_CAPACITY: usize = 4; + const DISPATCH_QUEUE_CAPACITY: usize = 4; + + let jobs = make_jobs(1); + 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( + 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); + reschedule_queue_sender + .send(injected) + .expect("reschedule queue sender closed"); + + let assignments = tick_and_drain_n(&mut scheduler, &reader, 1).await?; + let redispatched = assignments[0]; + assert_eq!(redispatched.job_id, job_a.0); + assert_eq!(redispatched.resource_group_id, job_a.1); + assert_eq!(redispatched.task_id, TaskId::Index(0)); + assert_eq!(redispatched.session_id, DEFAULT_SESSION_ID); + assert_ne!(redispatched.id, injected.id); + + assert_no_further_assignments(&mut scheduler, &reader).await?; + + Ok(()) +} + +#[tokio::test] +async fn reschedule_drops_stale_session_assignment() -> anyhow::Result<()> { + const ACTIVE_JOB_QUEUE_CAPACITY: usize = 4; + const DISPATCH_QUEUE_CAPACITY: usize = 4; + const NEW_SESSION_ID: SessionId = DEFAULT_SESSION_ID + 1; + + 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( + 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 + // assignment's. + tick_until(&mut scheduler, |scheduler| { + scheduler.storage_session_id == NEW_SESSION_ID + }) + .await?; + + let jobs = make_jobs(1); + let job_stale = jobs[0]; + let injected = make_reschedule_assignment(job_stale, TaskId::Index(0), DEFAULT_SESSION_ID, 0); + reschedule_queue_sender + .send(injected) + .expect("reschedule queue sender closed"); + + scheduler.tick().await?; + + assert!(!scheduler.active_jobs.contains_key(&job_stale.0)); + assert!( + scheduler + .buffered_tasks + .iter() + .all(|&(job_id, _)| job_id != job_stale.0), + "a stale-session assignment leaked into the buffered tasks", + ); + assert_no_further_assignments(&mut scheduler, &reader).await?; + + Ok(()) +} + +#[tokio::test] +async fn reschedule_dedups_against_inbound_ready_task() -> anyhow::Result<()> { + const ACTIVE_JOB_QUEUE_CAPACITY: usize = 4; + const DISPATCH_QUEUE_CAPACITY: usize = 1; + + let jobs = make_jobs(2); + let job_blocker = jobs[0]; + let job_a = jobs[1]; + + 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( + make_config(ACTIVE_JOB_QUEUE_CAPACITY, DISPATCH_QUEUE_CAPACITY), + storage_client, + writer, + ); + + tick_until(&mut scheduler, |scheduler| { + scheduler + .buffered_tasks + .contains(&(job_a.0, TaskId::Index(0))) + }) + .await?; + + let injected = make_reschedule_assignment(job_a, TaskId::Index(0), DEFAULT_SESSION_ID, 0); + reschedule_queue_sender + .send(injected) + .expect("reschedule queue sender closed"); + scheduler.tick().await?; + + // Unfreeze: exactly the blocker's task and `job_a`'s single task drain, proving the inbound and + // rescheduled copies of `job_a.t0` were deduplicated to a single dispatch. + let assignments = tick_and_drain_n(&mut scheduler, &reader, 2).await?; + assert_eq!( + make_assignment_tuple(&assignments), + vec![ + (job_blocker.0, job_blocker.1, TaskId::Index(0)), + (job_a.0, job_a.1, TaskId::Index(0)), + ], + ); + assert_no_further_assignments(&mut scheduler, &reader).await?; + + Ok(()) +} + +#[tokio::test(flavor = "multi_thread")] +async fn randomly_rescheduled_assignments_are_eventually_redispatched() -> anyhow::Result<()> { + const NUM_JOBS: usize = 4; + const TASKS_PER_JOB: usize = 5; + const DISPATCH_QUEUE_CAPACITY: usize = 32; + + let jobs = make_jobs(NUM_JOBS); + let storage_client = MockStorageClient::new(DEFAULT_SESSION_ID); + storage_client.push_ready_batch( + DEFAULT_SESSION_ID, + 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 total = NUM_JOBS * TASKS_PER_JOB; + let mut completed: HashSet<(JobId, TaskId)> = HashSet::new(); + let mut rescheduled: HashSet<(JobId, TaskId)> = HashSet::new(); + let mut seen_ids: HashSet = HashSet::new(); + let mut dispatch_count: HashMap<(JobId, TaskId), usize> = HashMap::new(); + + let deadline = tokio::time::Instant::now() + DRAIN_DEADLINE; + while completed.len() < total { + if tokio::time::Instant::now() > deadline { + bail!( + "timed out redispatching rescheduled assignments: completed {}, expected {total}", + completed.len(), + ); + } + let Some(assignment) = reader.dequeue(Duration::from_millis(100)).await? else { + continue; + }; + assert_eq!(assignment.session_id, DEFAULT_SESSION_ID); + assert!( + seen_ids.insert(assignment.id), + "a redispatched assignment reused an id: {:?}", + assignment.id, + ); + + let key = (assignment.job_id, assignment.task_id); + *dispatch_count.entry(key).or_insert(0) += 1; + + // Randomly decide whether to reschedule this assignment, but only if it has not already + // been rescheduled, so each task is rescheduled at most once. + if !rescheduled.contains(&key) && rand::random::() { + reschedule_sender + .send(assignment) + .expect("reschedule queue sender closed"); + rescheduled.insert(key); + } else { + completed.insert(key); + } + } + + let expected: HashSet<(JobId, TaskId)> = jobs + .iter() + .flat_map(|&(job_id, _)| (0..TASKS_PER_JOB).map(move |i| (job_id, TaskId::Index(i)))) + .collect(); + assert_eq!(completed, expected); + + // Every dispatch is accounted for: the totals add up to each task's original dispatch plus one + // extra dispatch per rescheduled task, and every rescheduled task is dispatched exactly twice. + let total_dispatches: usize = dispatch_count.values().sum(); + assert_eq!(total_dispatches, total + rescheduled.len()); + for key in &rescheduled { + assert_eq!( + dispatch_count.get(key).copied(), + Some(2), + "rescheduled task {key:?} was not dispatched exactly twice", + ); + } + + cancellation_token.cancel(); + scheduler_handle.await.expect("scheduler task panicked")?; + Ok(()) +} diff --git a/components/spider-scheduler/src/dispatch_queue.rs b/components/spider-scheduler/src/dispatch_queue.rs index 9f0e61431..23ccab176 100644 --- a/components/spider-scheduler/src/dispatch_queue.rs +++ b/components/spider-scheduler/src/dispatch_queue.rs @@ -59,21 +59,15 @@ pub trait DispatchQueueSource: Send + Sync + Clone { /// /// # Returns /// - /// `None` if no task assignment is available within the specified wait time, or a tuple - /// containing: - /// - /// * The storage session associated with the assignment. - /// * The next task assignment ready to execute. + /// * The next task assignment ready to execute, whose `session_id` carries the storage session. + /// * `None` if no task assignment is available within the specified wait time. /// /// # Errors /// /// Returns an error if: /// /// * [`SchedulerError::DispatchQueueClosed`] if the dispatching queue is closed. - async fn dequeue( - &self, - wait_time: Duration, - ) -> Result, SchedulerError>; + async fn dequeue(&self, wait_time: Duration) -> Result, SchedulerError>; } /// A cloneable writer handle for the dispatching queue, implementing [`DispatchQueueSink`] using @@ -128,16 +122,13 @@ pub struct DispatchQueueReader { #[async_trait] impl DispatchQueueSource for DispatchQueueReader { - async fn dequeue( - &self, - wait_time: Duration, - ) -> Result, SchedulerError> { + 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.inner.session_id.read().await; if let Ok(assignment) = self.inner.assignment_receiver.try_recv() { - return Ok(Some((*session_id_guard, assignment))); + return Ok(Some(assignment)); } if wait_time.is_zero() { @@ -145,7 +136,7 @@ impl DispatchQueueSource for DispatchQueueReader { } match tokio::time::timeout(wait_time, self.inner.assignment_receiver.recv()).await { - Ok(Ok(assignment)) => Ok(Some((*session_id_guard, assignment))), + Ok(Ok(assignment)) => Ok(Some(assignment)), Ok(Err(_)) => Err(SchedulerError::DispatchQueueClosed), Err(_) => Ok(None), } @@ -233,19 +224,21 @@ mod tests { /// /// Forwards [`make_assignment_with_task_id`]'s return values with `task_id` set with /// [`next_task_id`]'s return value. - fn make_assignment() -> TaskAssignment { - make_assignment_with_task_id(next_task_id()) + fn make_assignment(session_id: SessionId) -> TaskAssignment { + make_assignment_with_task_id(next_task_id(), session_id) } /// # Returns /// - /// A new [`TaskAssignment`] with the given `task_id` and other ID fields are auto-generated. - fn make_assignment_with_task_id(task_id: TaskId) -> TaskAssignment { + /// A new [`TaskAssignment`] with the given `task_id` and `session_id`; other ID fields are + /// auto-generated. + fn make_assignment_with_task_id(task_id: TaskId, session_id: SessionId) -> TaskAssignment { TaskAssignment { id: TaskAssignmentId::random(), resource_group_id: ResourceGroupId::random(), job_id: JobId::random(), task_id, + session_id, } } @@ -318,7 +311,7 @@ mod tests { let id = next_task_id(); tagged_for_writer.insert(id, current_session); writer - .enqueue(make_assignment_with_task_id(id)) + .enqueue(make_assignment_with_task_id(id, current_session)) .await .expect("enqueue failed"); } @@ -334,7 +327,7 @@ mod tests { let id = next_task_id(); tagged_for_writer.insert(id, current_session); writer - .enqueue(make_assignment_with_task_id(id)) + .enqueue(make_assignment_with_task_id(id, current_session)) .await .expect("enqueue failed"); } @@ -351,9 +344,9 @@ mod tests { tracker.spawn(async move { loop { match r.dequeue(Duration::from_millis(500)).await { - Ok(Some((session, assignment))) => { + Ok(Some(assignment)) => { if delivered_for_reader - .insert(assignment.task_id, session) + .insert(assignment.task_id, assignment.session_id) .is_some() { duplicates_for_reader.insert(assignment.task_id); @@ -401,15 +394,15 @@ mod tests { async fn sanity_round_trip_and_initial_session() -> Result<()> { const SESSION_ID: SessionId = 1; let (writer, reader) = create_dispatch_queue(8, SESSION_ID); - let assignment = make_assignment(); + let assignment = make_assignment(SESSION_ID); writer.enqueue(assignment).await?; - let (session, received) = reader + let received = reader .dequeue(Duration::from_millis(1)) .await? .expect("expected an assignment"); - assert_eq!(session, SESSION_ID); + assert_eq!(received.session_id, SESSION_ID); assert_eq!(received, assignment); Ok(()) } @@ -425,7 +418,7 @@ mod tests { for _ in 0..N { writer - .enqueue(make_assignment()) + .enqueue(make_assignment(1)) .await .expect("enqueue failed"); } @@ -450,7 +443,7 @@ mod tests { for _ in 0..N { writer - .enqueue(make_assignment()) + .enqueue(make_assignment(1)) .await .expect("enqueue failed"); tokio::time::sleep(Duration::from_millis(10)).await; @@ -498,21 +491,21 @@ mod tests { let (writer, reader) = create_dispatch_queue(8, SESSION_ID); writer.bump_session_id(NEW_SESSION_ID).await?; - writer.enqueue(make_assignment()).await?; + writer.enqueue(make_assignment(NEW_SESSION_ID)).await?; - let (session, _) = reader + let received = reader .dequeue(Duration::from_secs(1)) .await? .expect("expected an assignment"); - assert_eq!(session, NEW_SESSION_ID); + assert_eq!(received.session_id, NEW_SESSION_ID); Ok(()) } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn pre_bump_items_not_delivered() -> Result<()> { let (writer, reader) = create_dispatch_queue(8, 1); - writer.enqueue(make_assignment()).await?; - writer.enqueue(make_assignment()).await?; + writer.enqueue(make_assignment(1)).await?; + writer.enqueue(make_assignment(1)).await?; writer.bump_session_id(2).await?; let result = reader.dequeue(Duration::from_millis(100)).await?; @@ -524,14 +517,14 @@ mod tests { async fn post_bump_items_paired_with_new_session() -> Result<()> { let (writer, reader) = create_dispatch_queue(8, 1); writer.bump_session_id(2).await?; - let assignment = make_assignment(); + let assignment = make_assignment(2); writer.enqueue(assignment).await?; - let (session, received) = reader + let received = reader .dequeue(Duration::from_secs(1)) .await? .expect("expected an assignment"); - assert_eq!(session, 2); + assert_eq!(received.session_id, 2); assert_eq!(received, assignment); Ok(()) } @@ -553,21 +546,21 @@ mod tests { "expected InvalidSessionId(2), got {smaller:?}", ); - writer.enqueue(make_assignment()).await?; - let (session, _) = reader + writer.enqueue(make_assignment(3)).await?; + let received = reader .dequeue(Duration::from_secs(1)) .await? .expect("expected an assignment"); - assert_eq!(session, 3); + assert_eq!(received.session_id, 3); Ok(()) } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn size_zero_after_bump() -> Result<()> { let (writer, _reader) = create_dispatch_queue(8, 1); - writer.enqueue(make_assignment()).await?; - writer.enqueue(make_assignment()).await?; - writer.enqueue(make_assignment()).await?; + writer.enqueue(make_assignment(1)).await?; + writer.enqueue(make_assignment(1)).await?; + writer.enqueue(make_assignment(1)).await?; assert_eq!(writer.size(), 3); writer.bump_session_id(2).await?; @@ -592,7 +585,7 @@ mod tests { let writer_handle = tokio::spawn(async move { for &id in &pre_bump_for_writer { writer - .enqueue(make_assignment_with_task_id(id)) + .enqueue(make_assignment_with_task_id(id, INIT_SESSION)) .await .expect("enqueue failed"); } @@ -608,7 +601,7 @@ mod tests { for &id in &post_bump_for_writer { writer - .enqueue(make_assignment_with_task_id(id)) + .enqueue(make_assignment_with_task_id(id, MID_SESSION)) .await .expect("enqueue failed"); } @@ -621,7 +614,7 @@ mod tests { .expect("bump to final session failed"); writer - .enqueue(make_assignment_with_task_id(final_id)) + .enqueue(make_assignment_with_task_id(final_id, FINAL_SESSION)) .await .expect("enqueue failed"); drop(writer); @@ -630,8 +623,8 @@ mod tests { let mut delivered: HashMap = HashMap::new(); loop { match reader.dequeue(Duration::from_millis(100)).await { - Ok(Some((session, assignment))) => { - let prior = delivered.insert(assignment.task_id, session); + Ok(Some(assignment)) => { + let prior = delivered.insert(assignment.task_id, assignment.session_id); assert_eq!( prior, None, "duplicate delivery for {:?}", diff --git a/components/spider-scheduler/src/execution_manager_registry.rs b/components/spider-scheduler/src/execution_manager_registry.rs index 0b51238f8..2dd5547e6 100644 --- a/components/spider-scheduler/src/execution_manager_registry.rs +++ b/components/spider-scheduler/src/execution_manager_registry.rs @@ -354,6 +354,7 @@ mod tests { use spider_core::types::id::JobId; use spider_core::types::id::ResourceGroupId; + use spider_core::types::id::SessionId; use spider_core::types::id::TaskId; use tokio::sync::mpsc::UnboundedReceiver; use tokio::sync::mpsc::{self}; @@ -436,6 +437,7 @@ mod tests { resource_group_id: ResourceGroupId::random(), job_id: JobId::random(), task_id: TaskId::Index(0), + session_id: SessionId::default(), } } diff --git a/components/spider-scheduler/src/grpc.rs b/components/spider-scheduler/src/grpc.rs index 4d519e664..5cbb196cc 100644 --- a/components/spider-scheduler/src/grpc.rs +++ b/components/spider-scheduler/src/grpc.rs @@ -2,7 +2,6 @@ use async_trait::async_trait; use spider_core::types::id::SchedulerId; -use spider_core::types::id::SessionId; use spider_core::types::scheduler::TaskAssignment; use spider_proto_rust::common; use spider_proto_rust::scheduler::NextTaskResponse; @@ -134,14 +133,12 @@ impl SchedulerService .await .map_err(|error| self.service_error_handler(error, "next_task"))?; - let response = match dispatched { - Some((session_id, assignment)) => { - make_next_task_response(assignment, self.inner.scheduler_id(), session_id) - } - None => NextTaskResponse { + let response = dispatched.map_or( + NextTaskResponse { result: Some(next_task_response::Result::NoTask(common::Void {})), }, - }; + |assignment| make_next_task_response(assignment, self.inner.scheduler_id()), + ); Ok(Response::new(response)) } @@ -182,12 +179,10 @@ impl SchedulerService /// # Returns /// -/// A [`NextTaskResponse`] carrying the given assignment, stamped with `scheduler_id` and paired -/// with `session_id`. +/// A [`NextTaskResponse`] carrying the given assignment, stamped with `scheduler_id`. fn make_next_task_response( assignment: TaskAssignment, scheduler_id: SchedulerId, - session_id: SessionId, ) -> NextTaskResponse { NextTaskResponse { result: Some(next_task_response::Result::Assignment( @@ -197,7 +192,7 @@ fn make_next_task_response( job_id: assignment.job_id.get(), task_id: Some(common::TaskId::from(assignment.task_id)), scheduler_id: scheduler_id.get(), - session_id, + session_id: assignment.session_id, }, )), } diff --git a/components/spider-scheduler/src/runtime.rs b/components/spider-scheduler/src/runtime.rs index 123f6c1b9..bb2331589 100644 --- a/components/spider-scheduler/src/runtime.rs +++ b/components/spider-scheduler/src/runtime.rs @@ -1,9 +1,10 @@ //! The scheduler runtime. //! //! This module registers the scheduler with the storage service, wires the scheduler core to a -//! freshly created dispatch queue, 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. +//! 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. use std::time::Duration; @@ -22,7 +23,6 @@ use crate::execution_manager_registry::ExecutionManagerRegistry; use crate::execution_manager_registry::ExecutionManagerRegistryConfig; use crate::service::SchedulerServiceState; use crate::storage_client::SchedulerStorageClient; -use crate::types::TaskAssignment; /// Runtime configuration for the scheduler service. #[derive(Clone, Debug, Deserialize)] @@ -48,7 +48,6 @@ pub struct RuntimeConfig { /// Runtime state for the scheduler service. pub struct Runtime { core_join_handle: tokio::task::JoinHandle>, - _reschedule_queue_receiver: tokio::sync::mpsc::UnboundedReceiver, cancellation_token: CancellationToken, stop_timeout: Duration, } @@ -92,7 +91,8 @@ 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, and starts the core's scheduling loop as a background coroutine. +/// dispatch queue, hands the core the reschedule queue reader, and starts the core's scheduling +/// loop as a background coroutine. /// /// # Type Parameters /// @@ -153,6 +153,7 @@ pub async fn create_runtime>, stop_timeout_sec: u64, ) -> Runtime { - let (_reschedule_queue_sender, reschedule_queue_receiver) = - tokio::sync::mpsc::unbounded_channel(); Runtime { core_join_handle: core_task, - _reschedule_queue_receiver: reschedule_queue_receiver, cancellation_token, stop_timeout: Duration::from_secs(stop_timeout_sec), } diff --git a/components/spider-scheduler/src/service.rs b/components/spider-scheduler/src/service.rs index e835b175a..36117b6d4 100644 --- a/components/spider-scheduler/src/service.rs +++ b/components/spider-scheduler/src/service.rs @@ -11,7 +11,6 @@ use std::time::Duration; use spider_core::types::id::ExecutionManagerId; use spider_core::types::id::SchedulerId; -use spider_core::types::id::SessionId; use spider_core::types::scheduler::TaskAssignmentRecord; use crate::dispatch_queue::DispatchQueueSource; @@ -71,9 +70,8 @@ impl /// /// # Returns /// - /// * A tuple on success, containing: - /// * The storage session the dispatch queue paired with the assignment. - /// * The task assignment handed to the execution manager. + /// * The task assignment handed to the execution manager on success, whose `session_id` carries + /// the storage session. /// * `None` if no assignment becomes available within `wait_time`. /// /// # Errors @@ -86,7 +84,7 @@ impl em_id: ExecutionManagerId, prev_assignment: Option, wait_time: Duration, - ) -> Result, SchedulerServiceError> { + ) -> Result, SchedulerServiceError> { if let Some(prev) = prev_assignment { // The previous assignment is handled in a fire-and-forget task. Errors are ignored but // logged for observability purposes. @@ -106,7 +104,7 @@ impl ); Ok(None) } - Some((session_id, assignment)) => { + Some(assignment) => { tracing::info!( scheduler_id = % self.scheduler_id(), em_id = % em_id, @@ -116,7 +114,7 @@ impl "Task dispatched to execution manager." ); self.inner.registry.assign(em_id, assignment).await; - Ok(Some((session_id, assignment))) + Ok(Some(assignment)) } } } @@ -285,7 +283,7 @@ mod tests { async fn dequeue( &self, _wait_time: Duration, - ) -> Result, SchedulerError> { + ) -> Result, SchedulerError> { // Atomically claim one slot: return None once the counter is exhausted, otherwise // decrement and synthesize a fresh assignment. let mut current = self.remaining.load(Ordering::Relaxed); @@ -299,7 +297,7 @@ mod tests { Ordering::Relaxed, Ordering::Relaxed, ) { - Ok(_) => return Ok(Some((SESSION_ID, make_assignment()))), + Ok(_) => return Ok(Some(make_assignment())), Err(actual) => current = actual, } } @@ -319,6 +317,7 @@ mod tests { resource_group_id: ResourceGroupId::from(0), job_id: JobId::from(0), task_id: TaskId::Index(0), + session_id: SESSION_ID, } } @@ -406,11 +405,11 @@ mod tests { assert_eq!(service.scheduler_id(), SchedulerId::from(SCHEDULER_ID)); - let (session_id, assignment) = service + let assignment = service .next_task(em_id, None, Duration::from_millis(1)) .await? .expect("an assignment should be dequeued"); - assert_eq!(session_id, SESSION_ID); + assert_eq!(assignment.session_id, SESSION_ID); // The assignment was recorded against the execution manager, so shutting it down without // acknowledging the assignment reschedules it. @@ -444,15 +443,15 @@ mod tests { let (service, mut reschedule_queue_receiver, _cancellation_token) = build_service(3); let em_id = ExecutionManagerId::from(EM_ID); - let (_, assignment_a) = service + let assignment_a = service .next_task(em_id, None, Duration::from_millis(1)) .await? .expect("the first assignment should be dequeued"); - let (_, assignment_b) = service + let assignment_b = service .next_task(em_id, None, Duration::from_millis(1)) .await? .expect("the second assignment should be dequeued"); - let (_, assignment_c) = service + let assignment_c = service .next_task(em_id, None, Duration::from_millis(1)) .await? .expect("the third assignment should be dequeued"); diff --git a/tests/huntsman/em-runtime/tests/test_runtime.rs b/tests/huntsman/em-runtime/tests/test_runtime.rs index 1faca885e..41069a601 100644 --- a/tests/huntsman/em-runtime/tests/test_runtime.rs +++ b/tests/huntsman/em-runtime/tests/test_runtime.rs @@ -60,9 +60,9 @@ fn assignment_with_session(session_id: u64) -> SchedulerResponse { resource_group_id: ResourceGroupId::random(), job_id: JobId::random(), task_id: TaskId::Index(0), + session_id, }, scheduler_id: SchedulerId::random(), - session_id, } }