From 719e604991f9a03464b8773b1120d84224b014f8 Mon Sep 17 00:00:00 2001 From: Jeadie Date: Tue, 23 Jun 2026 13:47:47 +1000 Subject: [PATCH 1/4] feat: add a broadcast channel for job state event notifications Adds a tokio::sync::broadcast Sender to SchedulerServer, subscribable via subscribe_job_updates(). On job state changes (queued/running/completed/ failed/cancelled) the scheduler broadcasts a JobStateEvent carrying the job id and new state. Ported from spiceai/datafusion-ballista#15 for upstreaming. --- ballista/scheduler/src/lib.rs | 1 + .../src/scheduler_server/job_state_event.rs | 150 ++++++++++++++++++ .../scheduler/src/scheduler_server/mod.rs | 51 ++++++ .../scheduler_server/query_stage_scheduler.rs | 37 +++++ 4 files changed, 239 insertions(+) create mode 100644 ballista/scheduler/src/scheduler_server/job_state_event.rs diff --git a/ballista/scheduler/src/lib.rs b/ballista/scheduler/src/lib.rs index 533711ecb0..e2259b5cfa 100644 --- a/ballista/scheduler/src/lib.rs +++ b/ballista/scheduler/src/lib.rs @@ -47,3 +47,4 @@ mod flight_proxy_service; pub mod test_utils; pub use scheduler_server::SessionBuilder; +pub use scheduler_server::job_state_event::{JobState, JobStateEvent}; diff --git a/ballista/scheduler/src/scheduler_server/job_state_event.rs b/ballista/scheduler/src/scheduler_server/job_state_event.rs new file mode 100644 index 0000000000..df6a3c043c --- /dev/null +++ b/ballista/scheduler/src/scheduler_server/job_state_event.rs @@ -0,0 +1,150 @@ +// Licensed to the Apache Software Foundation (ASF) under one +// or more contributor license agreements. See the NOTICE file +// distributed with this work for additional information +// regarding copyright ownership. The ASF licenses this file +// to you under the Apache License, Version 2.0 (the +// "License"); you may not use this file except in compliance +// with the License. You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, +// software distributed under the License is distributed on an +// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY +// KIND, either express or implied. See the License for the +// specific language governing permissions and limitations +// under the License. + +//! Job state event notifications for subscribers. +//! +//! This module provides a broadcast-based notification system for job state changes. +//! Consumers can subscribe to receive notifications when jobs change state, avoiding +//! the need to poll for status updates. + +use std::fmt; + +/// Represents the current state of a job in the scheduler. +/// +/// This enum mirrors the possible states from the protobuf `job_status::Status` +/// but is designed to be lightweight for broadcasting. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum JobState { + /// Job is queued and waiting to be scheduled. + Queued, + /// Job is currently running. + Running, + /// Job has completed successfully. + Completed, + /// Job has failed with an error message. + Failed(String), + /// Job was cancelled. + Cancelled, +} + +impl fmt::Display for JobState { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + match self { + JobState::Queued => write!(f, "Queued"), + JobState::Running => write!(f, "Running"), + JobState::Completed => write!(f, "Completed"), + JobState::Failed(msg) => write!(f, "Failed: {}", msg), + JobState::Cancelled => write!(f, "Cancelled"), + } + } +} + +/// Event emitted when a job's state changes. +/// +/// This struct is designed to be cloned and sent through a broadcast channel +/// to notify subscribers about job state changes. +#[derive(Debug, Clone)] +pub struct JobStateEvent { + /// The unique identifier of the job. + pub job_id: String, + /// The new state of the job. + pub state: JobState, +} + +impl JobStateEvent { + /// Creates a new job state event. + pub fn new(job_id: impl Into, state: JobState) -> Self { + Self { + job_id: job_id.into(), + state, + } + } + + /// Creates a queued event for the given job. + pub fn queued(job_id: impl Into) -> Self { + Self::new(job_id, JobState::Queued) + } + + /// Creates a running event for the given job. + pub fn running(job_id: impl Into) -> Self { + Self::new(job_id, JobState::Running) + } + + /// Creates a completed event for the given job. + pub fn completed(job_id: impl Into) -> Self { + Self::new(job_id, JobState::Completed) + } + + /// Creates a failed event for the given job. + pub fn failed(job_id: impl Into, error: impl Into) -> Self { + Self::new(job_id, JobState::Failed(error.into())) + } + + /// Creates a cancelled event for the given job. + pub fn cancelled(job_id: impl Into) -> Self { + Self::new(job_id, JobState::Cancelled) + } +} + +impl fmt::Display for JobStateEvent { + fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result { + write!( + f, + "JobStateEvent[job_id={}, state={}]", + self.job_id, self.state + ) + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_job_state_event_creation() { + let event = JobStateEvent::queued("job-123"); + assert_eq!(event.job_id, "job-123"); + assert_eq!(event.state, JobState::Queued); + + let event = JobStateEvent::running("job-123"); + assert_eq!(event.state, JobState::Running); + + let event = JobStateEvent::completed("job-123"); + assert_eq!(event.state, JobState::Completed); + + let event = JobStateEvent::failed("job-123", "Something went wrong"); + assert_eq!( + event.state, + JobState::Failed("Something went wrong".to_string()) + ); + + let event = JobStateEvent::cancelled("job-123"); + assert_eq!(event.state, JobState::Cancelled); + } + + #[test] + fn test_job_state_display() { + assert_eq!(JobState::Queued.to_string(), "Queued"); + assert_eq!(JobState::Running.to_string(), "Running"); + assert_eq!(JobState::Completed.to_string(), "Completed"); + assert_eq!( + JobState::Failed("error".to_string()).to_string(), + "Failed: error" + ); + assert_eq!(JobState::Cancelled.to_string(), "Cancelled"); + } +} diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index 17b955a98c..dc5e060eb7 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -23,6 +23,7 @@ use ballista_core::event_loop::{EventLoop, EventSender}; use ballista_core::serde::BallistaCodec; use ballista_core::serde::protobuf::TaskStatus; use ballista_core::{JobId, JobStatusSubscriber}; +use tokio::sync::broadcast; use datafusion::execution::context::SessionState; use datafusion::logical_expr::LogicalPlan; @@ -57,6 +58,8 @@ pub mod event; #[cfg(feature = "keda-scaler")] mod external_scaler; mod grpc; +/// Job state event notifications for subscribers. +pub mod job_state_event; pub(crate) mod query_stage_scheduler; /// Function type for building DataFusion session states from configuration. @@ -85,9 +88,20 @@ pub struct SchedulerServer>, /// Scheduler configuration. config: Arc, + /// Broadcast sender for job state change notifications. + /// + /// Subscribers can receive notifications when jobs change state by calling + /// `subscribe_job_updates()`. + job_state_sender: broadcast::Sender, } impl SchedulerServer { + /// Default capacity for the job state broadcast channel. + /// + /// This determines how many job state events can be buffered before + /// slow receivers start lagging behind. + const JOB_STATE_CHANNEL_CAPACITY: usize = 256; + /// Creates a new `SchedulerServer` with the given configuration. pub fn new( scheduler_name: String, @@ -102,10 +116,12 @@ impl SchedulerServer SchedulerServer SchedulerServer SchedulerServer SchedulerServer usize { self.state.task_manager.running_job_number() } + + /// Subscribes to job state change notifications. + /// + /// Returns a receiver that will receive [`JobStateEvent`] notifications + /// whenever a job changes state. This allows consumers to be notified + /// of job state changes without polling. + /// + /// # Example + /// + /// ```ignore + /// let mut receiver = scheduler.subscribe_job_updates(); + /// tokio::spawn(async move { + /// while let Ok(event) = receiver.recv().await { + /// println!("Job {} changed to state: {}", event.job_id, event.state); + /// } + /// }); + /// ``` + /// + /// # Note + /// + /// If the receiver falls behind and the channel buffer fills up, + /// older messages will be dropped and the receiver will receive + /// a `RecvError::Lagged` error on the next `recv()` call. + /// + /// [`JobStateEvent`]: job_state_event::JobStateEvent + pub fn subscribe_job_updates( + &self, + ) -> broadcast::Receiver { + self.job_state_sender.subscribe() + } + #[cfg(feature = "rest-api")] pub(crate) fn metrics_collector(&self) -> &dyn SchedulerMetricsCollector { self.query_stage_scheduler.metrics_collector() diff --git a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs index 0d84950bc7..0fb0898bdd 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -20,6 +20,7 @@ use std::time::Duration; use ballista_core::serde::protobuf::{FailedJob, JobStatus}; use log::{debug, error, info, trace, warn}; +use tokio::sync::broadcast; use ballista_core::error::{BallistaError, Result}; use ballista_core::event_loop::{EventAction, EventSender}; @@ -27,6 +28,7 @@ use tokio::sync::mpsc::error::TrySendError; use crate::config::SchedulerConfig; use crate::metrics::SchedulerMetricsCollector; +use crate::scheduler_server::job_state_event::JobStateEvent; use crate::scheduler_server::timestamp_millis; use datafusion_proto::logical_plan::AsLogicalPlan; use datafusion_proto::physical_plan::AsExecutionPlan; @@ -44,6 +46,8 @@ pub(crate) struct QueryStageScheduler< state: Arc>, metrics_collector: Arc, config: Arc, + /// Broadcast sender for job state change notifications. + job_state_sender: broadcast::Sender, } impl QueryStageScheduler { @@ -51,13 +55,22 @@ impl QueryStageSchedul state: Arc>, metrics_collector: Arc, config: Arc, + job_state_sender: broadcast::Sender, ) -> Self { Self { state, metrics_collector, config, + job_state_sender, } } + + /// Broadcasts a job state event to all subscribers. + fn broadcast_job_state(&self, event: JobStateEvent) { + // Ignore send errors - no receivers is a valid state + let _ = self.job_state_sender.send(event); + } + #[cfg(feature = "rest-api")] pub(crate) fn metrics_collector(&self) -> &dyn SchedulerMetricsCollector { self.metrics_collector.as_ref() @@ -98,6 +111,9 @@ impl } => { info!("Job queued: [{job_id}]"); + // Broadcast job queued state + self.broadcast_job_state(JobStateEvent::queued(&job_id)); + if let Err(e) = self .state .task_manager @@ -108,6 +124,9 @@ impl } let state = self.state.clone(); + + // Clone the job state sender to move into the async task + let job_state_sender = self.job_state_sender.clone(); tokio::spawn(async move { let event = if let Err(e) = state .submit_job( @@ -154,6 +173,8 @@ impl failed_at: timestamp_millis(), } } else { + // Broadcast job running state when successfully submitted + let _ = job_state_sender.send(JobStateEvent::running(&job_id)); QueryStageSchedulerEvent::JobSubmitted { job_id, queued_at, @@ -191,6 +212,10 @@ impl .record_failed(&job_id, queued_at, failed_at); error!("Job {job_id} failed: {fail_message}"); + + // Broadcast job failed state + self.broadcast_job_state(JobStateEvent::failed(&job_id, &fail_message)); + if let Err(e) = self .state .task_manager @@ -211,6 +236,10 @@ impl .record_completed(&job_id, queued_at, completed_at); info!("Job finished successfully: [{job_id}]"); + + // Broadcast job completed state + self.broadcast_job_state(JobStateEvent::completed(&job_id)); + if let Err(e) = self.state.task_manager.succeed_job(&job_id).await { error!("Fail to invoke succeed_job for job {job_id} due to {e:?}"); } @@ -226,6 +255,10 @@ impl .record_failed(&job_id, queued_at, failed_at); error!("Job failed: [{job_id}]"); + + // Broadcast job failed state + self.broadcast_job_state(JobStateEvent::failed(&job_id, &fail_message)); + match self .state .task_manager @@ -257,6 +290,10 @@ impl self.metrics_collector.record_cancelled(&job_id); info!("Job cancelled: [{job_id}]"); + + // Broadcast job cancelled state + self.broadcast_job_state(JobStateEvent::cancelled(&job_id)); + match self.state.task_manager.cancel_job(&job_id).await { Ok((running_tasks, _pending_tasks)) => { event_sender From 6ed1d3e4b27a62c37e49e54699afbc4143e657e2 Mon Sep 17 00:00:00 2001 From: jeadie Date: Tue, 23 Jun 2026 15:20:54 +1000 Subject: [PATCH 2/4] fix(core): add From<&JobId> for String so borrowed ids satisfy Into --- ballista/core/src/ids.rs | 6 ++++++ 1 file changed, 6 insertions(+) diff --git a/ballista/core/src/ids.rs b/ballista/core/src/ids.rs index c10d23b376..1df1a6214b 100644 --- a/ballista/core/src/ids.rs +++ b/ballista/core/src/ids.rs @@ -99,6 +99,12 @@ macro_rules! string_id { } } + impl From<&$name> for String { + fn from(value: &$name) -> Self { + value.0.clone() + } + } + impl AsRef for $name { fn as_ref(&self) -> &str { &self.0 From 1323db6c9587d6f5b6a8ef2181c295780f517e17 Mon Sep 17 00:00:00 2001 From: jeadie Date: Fri, 3 Jul 2026 11:42:29 +1000 Subject: [PATCH 3/4] fix(scheduler): address PR review feedback on job state broadcast channel - Add job_state_channel_capacity to SchedulerConfig (default 256) so operators with high job throughput can tune the buffer; removes the hardcoded const from SchedulerServer - Move all job state broadcasts to after their state transitions succeed, so subscribers never see an event for a commit that failed; aligns Queued/Failed/Completed/Cancelled with how Running already behaved - Add From for JobState to keep the two enums in sync; document that Cancelled has no protobuf counterpart - Expand module doc with guidance on when to use JobStateEvent vs JobStatusSubscriber, and reference job_state_channel_capacity for tuning - Add integration tests covering the broadcast wiring: lifecycle order (Queued -> Running -> Completed) and cancel path --- ballista/scheduler/src/config.rs | 13 ++ .../src/scheduler_server/job_state_event.rs | 41 +++++- .../scheduler/src/scheduler_server/mod.rs | 10 +- .../scheduler_server/query_stage_scheduler.rs | 129 +++++++++++++++--- ballista/scheduler/src/test_utils.rs | 9 ++ 5 files changed, 174 insertions(+), 28 deletions(-) diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index 3c6cdb0074..d1256044e8 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -269,6 +269,11 @@ pub struct SchedulerConfig { pub task_max_failures: usize, /// Number of failures attempts before stage is considered failed pub stage_max_failures: usize, + /// Capacity of the job state broadcast channel. + /// + /// Controls how many job state events can be buffered before slow receivers + /// start lagging. Increase this for schedulers with high job throughput. + pub job_state_channel_capacity: usize, #[cfg(feature = "rest-api")] /// Should the rest api be disabled pub disable_rest_api: bool, @@ -308,6 +313,7 @@ impl Default for SchedulerConfig { use_tls: false, task_max_failures: 4, stage_max_failures: 4, + job_state_channel_capacity: 256, #[cfg(feature = "rest-api")] disable_rest_api: false, #[cfg(feature = "rest-api")] @@ -449,6 +455,12 @@ impl SchedulerConfig { self.use_tls = use_tls; self } + + /// Sets the capacity of the job state broadcast channel. + pub fn with_job_state_channel_capacity(mut self, capacity: usize) -> Self { + self.job_state_channel_capacity = capacity; + self + } } /// Policy of distributing tasks to available executor slots @@ -540,6 +552,7 @@ impl TryFrom for SchedulerConfig { use_tls: false, task_max_failures: opt.task_max_failures, stage_max_failures: opt.stage_max_failures, + job_state_channel_capacity: 256, #[cfg(feature = "rest-api")] disable_rest_api: opt.disable_rest_api, #[cfg(feature = "rest-api")] diff --git a/ballista/scheduler/src/scheduler_server/job_state_event.rs b/ballista/scheduler/src/scheduler_server/job_state_event.rs index df6a3c043c..a5476791e8 100644 --- a/ballista/scheduler/src/scheduler_server/job_state_event.rs +++ b/ballista/scheduler/src/scheduler_server/job_state_event.rs @@ -20,13 +20,36 @@ //! This module provides a broadcast-based notification system for job state changes. //! Consumers can subscribe to receive notifications when jobs change state, avoiding //! the need to poll for status updates. - +//! +//! # Choosing between `JobStateEvent` and `JobStatusSubscriber` +//! +//! Ballista has two mechanisms for observing job progress: +//! +//! - **[`JobStateEvent`] / [`crate::scheduler_server::SchedulerServer::subscribe_job_updates`]**: +//! a `tokio::sync::broadcast` channel that delivers lifecycle events for *all* jobs +//! on this scheduler. Suited for cluster-wide observers such as metrics collectors, +//! audit logs, or HA state replication. Subscribers receive a lightweight event +//! containing only the job ID and new state; they must query the scheduler separately +//! for full job details. Slow subscribers may lag and miss events if the channel buffer fills; +//! see [`crate::config::SchedulerConfig::job_state_channel_capacity`] to tune this. +//! +//! - **`JobStatusSubscriber`** (`tokio::sync::mpsc::Sender`): a per-job +//! `mpsc` channel threaded through [`crate::scheduler_server::SchedulerServer::submit_job`]. +//! Suited for a single caller that submitted a job and wants rich status updates +//! (including partition locations on success) for that specific job. Backpressure +//! is applied to the scheduler if the subscriber is slow, so this is best used by +//! in-process callers that consume updates promptly. + +use ballista_core::serde::protobuf::job_status; use std::fmt; /// Represents the current state of a job in the scheduler. /// -/// This enum mirrors the possible states from the protobuf `job_status::Status` -/// but is designed to be lightweight for broadcasting. +/// Mirrors the variants of the protobuf `job_status::Status` but is +/// lightweight for broadcasting. When adding or renaming variants here, +/// update the `From` impl below to keep them in sync. +/// Note that `Cancelled` has no protobuf counterpart; cancelled jobs are +/// represented as `Failed` in the wire format. #[derive(Debug, Clone, PartialEq, Eq)] pub enum JobState { /// Job is queued and waiting to be scheduled. @@ -53,6 +76,18 @@ impl fmt::Display for JobState { } } +impl From for JobState { + fn from(status: job_status::Status) -> Self { + match status { + job_status::Status::Queued(_) => JobState::Queued, + job_status::Status::Running(_) => JobState::Running, + job_status::Status::Failed(f) => JobState::Failed(f.error), + // Protobuf has no Cancelled variant; cancelled jobs surface as Failed on the wire. + job_status::Status::Successful(_) => JobState::Completed, + } + } +} + /// Event emitted when a job's state changes. /// /// This struct is designed to be cloned and sent through a broadcast channel diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index dc5e060eb7..7a03276ff2 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -96,12 +96,6 @@ pub struct SchedulerServer SchedulerServer { - /// Default capacity for the job state broadcast channel. - /// - /// This determines how many job state events can be buffered before - /// slow receivers start lagging behind. - const JOB_STATE_CHANNEL_CAPACITY: usize = 256; - /// Creates a new `SchedulerServer` with the given configuration. pub fn new( scheduler_name: String, @@ -116,7 +110,7 @@ impl SchedulerServer SchedulerServer } => { info!("Job queued: [{job_id}]"); - // Broadcast job queued state - self.broadcast_job_state(JobStateEvent::queued(&job_id)); - if let Err(e) = self .state .task_manager @@ -123,6 +120,9 @@ impl return Ok(()); } + // Broadcast after state transition succeeds + self.broadcast_job_state(JobStateEvent::queued(&job_id)); + let state = self.state.clone(); // Clone the job state sender to move into the async task @@ -213,19 +213,19 @@ impl error!("Job {job_id} failed: {fail_message}"); - // Broadcast job failed state - self.broadcast_job_state(JobStateEvent::failed(&job_id, &fail_message)); - if let Err(e) = self .state .task_manager - .fail_unscheduled_job(&job_id, fail_message) + .fail_unscheduled_job(&job_id, fail_message.clone()) .await { error!( "Fail to invoke fail_unscheduled_job for job {job_id} due to {e:?}" ); } + + // Broadcast after state transition + self.broadcast_job_state(JobStateEvent::failed(&job_id, &fail_message)); } QueryStageSchedulerEvent::JobFinished { job_id, @@ -237,11 +237,11 @@ impl info!("Job finished successfully: [{job_id}]"); - // Broadcast job completed state - self.broadcast_job_state(JobStateEvent::completed(&job_id)); - if let Err(e) = self.state.task_manager.succeed_job(&job_id).await { error!("Fail to invoke succeed_job for job {job_id} due to {e:?}"); + } else { + // Broadcast only after state transition succeeds + self.broadcast_job_state(JobStateEvent::completed(&job_id)); } self.state.clean_up_successful_job(job_id); } @@ -256,16 +256,15 @@ impl error!("Job failed: [{job_id}]"); - // Broadcast job failed state - self.broadcast_job_state(JobStateEvent::failed(&job_id, &fail_message)); - match self .state .task_manager - .abort_job(&job_id, fail_message) + .abort_job(&job_id, fail_message.clone()) .await { Ok((running_tasks, _pending_tasks)) => { + // Broadcast only after state transition succeeds + self.broadcast_job_state(JobStateEvent::failed(&job_id, &fail_message)); if !running_tasks.is_empty() { event_sender .post_event(QueryStageSchedulerEvent::CancelTasks( @@ -291,11 +290,10 @@ impl info!("Job cancelled: [{job_id}]"); - // Broadcast job cancelled state - self.broadcast_job_state(JobStateEvent::cancelled(&job_id)); - match self.state.task_manager.cancel_job(&job_id).await { Ok((running_tasks, _pending_tasks)) => { + // Broadcast only after state transition succeeds + self.broadcast_job_state(JobStateEvent::cancelled(&job_id)); event_sender .post_event(QueryStageSchedulerEvent::CancelTasks( running_tasks, @@ -404,6 +402,7 @@ impl #[cfg(test)] mod tests { use crate::config::SchedulerConfig; + use crate::scheduler_server::job_state_event::{JobState, JobStateEvent}; use crate::test_utils::{SchedulerTest, TestMetricsCollector, await_condition}; use ballista_core::config::TaskSchedulingPolicy; use ballista_core::error::Result; @@ -468,6 +467,102 @@ mod tests { Ok(()) } + /// Collects events from a broadcast receiver until `predicate` returns true or timeout elapses. + async fn collect_events_until( + rx: &mut tokio::sync::broadcast::Receiver, + predicate: impl Fn(&[JobState]) -> bool, + timeout: Duration, + ) -> Vec { + let mut states = vec![]; + let deadline = tokio::time::Instant::now() + timeout; + loop { + let remaining = deadline.saturating_duration_since(tokio::time::Instant::now()); + if remaining.is_zero() { + break; + } + match tokio::time::timeout(remaining, rx.recv()).await { + Ok(Ok(e)) => { + states.push(e.state.clone()); + if predicate(&states) { + break; + } + } + _ => break, + } + } + states + } + + #[tokio::test] + async fn test_job_state_broadcast_lifecycle_order() -> Result<()> { + let plan = test_plan(1); + let metrics_collector = Arc::new(TestMetricsCollector::default()); + let mut test = SchedulerTest::new( + SchedulerConfig::default() + .with_scheduler_policy(TaskSchedulingPolicy::PushStaged), + metrics_collector, + 1, + 1, + None, + ) + .await?; + + let mut rx = test.subscribe_job_updates(); + test.run("", &plan).await?; + + let states = collect_events_until( + &mut rx, + |s| s.contains(&JobState::Completed), + Duration::from_secs(5), + ) + .await; + + let queued_pos = states.iter().position(|s| *s == JobState::Queued) + .expect("Expected Queued event"); + let running_pos = states.iter().position(|s| *s == JobState::Running) + .expect("Expected Running event"); + let completed_pos = states.iter().position(|s| *s == JobState::Completed) + .expect("Expected Completed event"); + assert!(queued_pos < running_pos, "Queued should precede Running, got: {states:?}"); + assert!(running_pos < completed_pos, "Running should precede Completed, got: {states:?}"); + + Ok(()) + } + + #[tokio::test] + async fn test_job_state_broadcast_cancel() -> Result<()> { + let plan = test_plan(10); + let metrics_collector = Arc::new(TestMetricsCollector::default()); + let mut test = SchedulerTest::new( + SchedulerConfig::default() + .with_scheduler_policy(TaskSchedulingPolicy::PushStaged), + metrics_collector, + 1, + 1, + None, + ) + .await?; + + let mut rx = test.subscribe_job_updates(); + let job_id = test.submit("", &plan).await?; + test.tick().await?; + test.cancel(&job_id).await?; + + let states = collect_events_until( + &mut rx, + |s| s.contains(&JobState::Cancelled), + Duration::from_secs(5), + ) + .await; + + assert!( + states.contains(&JobState::Cancelled), + "Expected Cancelled event, got: {states:?}" + ); + + Ok(()) + } + fn test_plan(partitions: usize) -> LogicalPlan { let schema = Schema::new(vec![ Field::new("id", DataType::Utf8, false), diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index 810983ef2b..5c1263c391 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -544,6 +544,15 @@ impl SchedulerTest { Ok(()) } + /// Subscribes to job state change events from the underlying scheduler. + pub fn subscribe_job_updates( + &self, + ) -> tokio::sync::broadcast::Receiver< + crate::scheduler_server::job_state_event::JobStateEvent, + > { + self.scheduler.subscribe_job_updates() + } + /// Cancels a job by ID. pub async fn cancel(&self, job_id: &JobId) -> Result<()> { self.scheduler From 102fce2c55a1f3a8dffbd7c31896760c1262b26b Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> Date: Mon, 20 Jul 2026 09:54:37 +0900 Subject: [PATCH 4/4] docs: link JobStateEvent by absolute path so rustdoc resolves it --- ballista/scheduler/src/scheduler_server/job_state_event.rs | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/ballista/scheduler/src/scheduler_server/job_state_event.rs b/ballista/scheduler/src/scheduler_server/job_state_event.rs index a5476791e8..612adab2c9 100644 --- a/ballista/scheduler/src/scheduler_server/job_state_event.rs +++ b/ballista/scheduler/src/scheduler_server/job_state_event.rs @@ -25,7 +25,8 @@ //! //! Ballista has two mechanisms for observing job progress: //! -//! - **[`JobStateEvent`] / [`crate::scheduler_server::SchedulerServer::subscribe_job_updates`]**: +//! - **[`JobStateEvent`](crate::scheduler_server::job_state_event::JobStateEvent) / +//! [`crate::scheduler_server::SchedulerServer::subscribe_job_updates`]**: //! a `tokio::sync::broadcast` channel that delivers lifecycle events for *all* jobs //! on this scheduler. Suited for cluster-wide observers such as metrics collectors, //! audit logs, or HA state replication. Subscribers receive a lightweight event