From 719e604991f9a03464b8773b1120d84224b014f8 Mon Sep 17 00:00:00 2001 From: Jeadie Date: Tue, 23 Jun 2026 13:47:47 +1000 Subject: [PATCH 1/6] 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/6] 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/6] 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/6] 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 From c4b8cba75970e771789cf61aae7c632b9dab3486 Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> Date: Wed, 5 Aug 2026 12:01:47 +0900 Subject: [PATCH 5/6] refactor(scheduler): expose existing job state event stream --- ballista/core/src/ids.rs | 6 - ballista/scheduler/src/cluster/memory.rs | 66 +++++-- ballista/scheduler/src/config.rs | 13 -- ballista/scheduler/src/lib.rs | 1 - .../src/scheduler_server/job_state_event.rs | 186 ------------------ .../scheduler/src/scheduler_server/mod.rs | 48 +---- .../scheduler_server/query_stage_scheduler.rs | 48 +++++ ballista/scheduler/src/state/task_manager.rs | 7 +- ballista/scheduler/src/test_utils.rs | 16 +- 9 files changed, 114 insertions(+), 277 deletions(-) delete mode 100644 ballista/scheduler/src/scheduler_server/job_state_event.rs diff --git a/ballista/core/src/ids.rs b/ballista/core/src/ids.rs index 1df1a6214b..c10d23b376 100644 --- a/ballista/core/src/ids.rs +++ b/ballista/core/src/ids.rs @@ -99,12 +99,6 @@ 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 diff --git a/ballista/scheduler/src/cluster/memory.rs b/ballista/scheduler/src/cluster/memory.rs index 51dfa37dce..3baef1dddb 100644 --- a/ballista/scheduler/src/cluster/memory.rs +++ b/ballista/scheduler/src/cluster/memory.rs @@ -493,22 +493,20 @@ impl JobState for InMemoryJobState { async fn fail_unscheduled_job(&self, job_id: &JobId, reason: String) -> Result<()> { if let Some((job_id, (job_name, queued_at))) = self.queued_jobs.remove(job_id) { - self.completed_jobs.insert( - job_id.clone(), - ( - JobStatus { - job_id: job_id.into(), - job_name, - status: Some(Status::Failed(FailedJob { - error: reason, - queued_at, - started_at: 0, - ended_at: timestamp_millis(), - })), - }, - None, - ), - ); + let status = JobStatus { + job_id: job_id.clone().into(), + job_name, + status: Some(Status::Failed(FailedJob { + error: reason, + queued_at, + started_at: 0, + ended_at: timestamp_millis(), + })), + }; + self.completed_jobs + .insert(job_id.clone(), (status.clone(), None)); + self.job_event_sender + .send(&JobStateEvent::JobUpdated { job_id, status }); Ok(()) } else { @@ -533,6 +531,7 @@ mod test { use crate::test_utils::{ test_aggregation_plan, test_join_plan, test_two_aggregations_plan, }; + use ballista_core::JobId; use ballista_core::error::Result; use ballista_core::serde::protobuf::JobStatus; use ballista_core::serde::scheduler::{ @@ -609,6 +608,41 @@ mod test { Ok(()) } + #[tokio::test] + async fn test_in_memory_job_planning_failure_notification() -> Result<()> { + let state = InMemoryJobState::new( + "", + Arc::new(default_session_builder), + Arc::new(default_config_producer), + ); + let mut events = state.job_state_events().await?; + let job_id = JobId::from("job-1"); + + state.accept_job(&job_id, "", 0)?; + state + .fail_unscheduled_job(&job_id, "failed planning".to_owned()) + .await?; + + let event = + tokio::time::timeout(std::time::Duration::from_secs(1), events.next()) + .await + .expect("job state event should arrive") + .expect("job state event stream should remain open"); + + assert!(matches!( + event, + JobStateEvent::JobUpdated { + job_id: event_job_id, + status: JobStatus { + status: Some(ballista_core::serde::protobuf::job_status::Status::Failed(_)), + .. + }, + } if event_job_id == job_id + )); + + Ok(()) + } + #[tokio::test] async fn test_in_memory_job_notification() -> Result<()> { let state = InMemoryJobState::new( diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index f199b10fbb..ff935374bb 100644 --- a/ballista/scheduler/src/config.rs +++ b/ballista/scheduler/src/config.rs @@ -304,11 +304,6 @@ 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, @@ -351,7 +346,6 @@ impl Default for SchedulerConfig { min_ready_executors: 1, 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")] @@ -506,12 +500,6 @@ 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 @@ -606,7 +594,6 @@ impl TryFrom for SchedulerConfig { min_ready_executors: opt.min_ready_executors, 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/lib.rs b/ballista/scheduler/src/lib.rs index e2259b5cfa..533711ecb0 100644 --- a/ballista/scheduler/src/lib.rs +++ b/ballista/scheduler/src/lib.rs @@ -47,4 +47,3 @@ 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 deleted file mode 100644 index 612adab2c9..0000000000 --- a/ballista/scheduler/src/scheduler_server/job_state_event.rs +++ /dev/null @@ -1,186 +0,0 @@ -// 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. -//! -//! # Choosing between `JobStateEvent` and `JobStatusSubscriber` -//! -//! Ballista has two mechanisms for observing job progress: -//! -//! - **[`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 -//! 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. -/// -/// 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. - 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"), - } - } -} - -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 -/// 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 db3af7594f..f16189aae1 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -23,7 +23,6 @@ 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; @@ -32,7 +31,7 @@ use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_proto::logical_plan::AsLogicalPlan; use datafusion_proto::physical_plan::AsExecutionPlan; -use crate::cluster::BallistaCluster; +use crate::cluster::{BallistaCluster, JobStateEventStream}; use crate::config::SchedulerConfig; use crate::metrics::SchedulerMetricsCollector; use ballista_core::serde::scheduler::{ExecutorData, ExecutorMetadata}; @@ -59,8 +58,6 @@ 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. @@ -89,11 +86,6 @@ 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 { @@ -111,12 +103,10 @@ impl SchedulerServer SchedulerServer SchedulerServer SchedulerServer SchedulerServer broadcast::Receiver { - self.job_state_sender.subscribe() + /// Returns a stream of job state events from the configured state backend. + pub async fn job_state_events(&self) -> Result { + self.state.task_manager.job_state_events().await } /// True when at least `min_ready_executors` executors currently have @@ -235,7 +196,6 @@ impl SchedulerServer= self.state.config.min_ready_executors } - #[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 4b64e205a6..8b9417d388 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -465,6 +465,7 @@ impl #[cfg(test)] mod tests { + use crate::cluster::JobStateEvent; use crate::config::SchedulerConfig; use crate::test_utils::{SchedulerTest, TestMetricsCollector, await_condition}; use ballista_core::config::TaskSchedulingPolicy; @@ -474,6 +475,7 @@ mod tests { use datafusion::functions_aggregate::sum::sum; use datafusion::logical_expr::{LogicalPlan, col}; use datafusion::test_util::scan_empty_with_partitions; + use futures::StreamExt; use std::sync::Arc; use std::time::Duration; @@ -531,6 +533,52 @@ mod tests { Ok(()) } + #[tokio::test] + async fn test_scheduler_exposes_job_state_events() -> 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 events = test.job_state_events().await?; + let (_, job_id) = test.run("", &plan).await?; + + let received_success = tokio::time::timeout(Duration::from_secs(5), async { + while let Some(event) = events.next().await { + if matches!( + event, + JobStateEvent::JobUpdated { + job_id: event_job_id, + status, + } if event_job_id == job_id + && matches!( + status.status, + Some(job_status::Status::Successful(_)) + ) + ) { + return true; + } + } + false + }) + .await + .expect("successful job state event should arrive"); + + assert!( + received_success, + "job state event stream closed unexpectedly" + ); + + Ok(()) + } + #[tokio::test] async fn test_running_job_fails_when_all_executors_are_lost() -> Result<()> { let plan = test_plan(10); diff --git a/ballista/scheduler/src/state/task_manager.rs b/ballista/scheduler/src/state/task_manager.rs index 7d46488c77..dea5af1004 100644 --- a/ballista/scheduler/src/state/task_manager.rs +++ b/ballista/scheduler/src/state/task_manager.rs @@ -15,7 +15,7 @@ // specific language governing permissions and limitations // under the License. -use crate::cluster::JobState; +use crate::cluster::{JobState, JobStateEventStream}; use crate::config::SchedulerConfig; use crate::planner::DefaultDistributedPlanner; use crate::scheduler_server::event::{QueryStageSchedulerEvent, SubmitPlan}; @@ -243,6 +243,11 @@ impl TaskManager .count() } + /// Returns a stream of job state events from the configured state backend. + pub async fn job_state_events(&self) -> Result { + self.state.job_state_events().await + } + /// Generate an ExecutionGraph for the job and save it to the persistent state. /// By default, this job will be curated by the scheduler which receives it. /// Then we will also save it to the active execution graph diff --git a/ballista/scheduler/src/test_utils.rs b/ballista/scheduler/src/test_utils.rs index 31f8de31a7..496aa8d587 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -53,7 +53,7 @@ use datafusion::physical_plan::display::DisplayableExecutionPlan; use datafusion::prelude::{CsvReadOptions, JoinType, col}; use datafusion::test_util::scan_empty_with_partitions; -use crate::cluster::BallistaCluster; +use crate::cluster::{BallistaCluster, JobStateEventStream}; use crate::scheduler_server::event::QueryStageSchedulerEvent; use crate::state::execution_graph::{ @@ -479,6 +479,11 @@ impl SchedulerTest { self.scheduler.running_job_number() } + /// Returns job state events from the underlying scheduler. + pub async fn job_state_events(&self) -> Result { + self.scheduler.job_state_events().await + } + /// Returns the session context for tests. pub async fn ctx(&self) -> Result> { self.scheduler @@ -534,15 +539,6 @@ 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 bd25aeddce66002f0aaef9f01a78b6adac459e24 Mon Sep 17 00:00:00 2001 From: Phillip LeBlanc <879445+phillipleblanc@users.noreply.github.com> Date: Wed, 5 Aug 2026 17:48:20 +0900 Subject: [PATCH 6/6] fix(scheduler): expose cluster state event stream --- .../scheduler/src/scheduler_server/mod.rs | 35 ++++++++++++++++++- .../scheduler/src/state/executor_manager.rs | 7 +++- 2 files changed, 40 insertions(+), 2 deletions(-) diff --git a/ballista/scheduler/src/scheduler_server/mod.rs b/ballista/scheduler/src/scheduler_server/mod.rs index f16189aae1..ca439dfeb2 100644 --- a/ballista/scheduler/src/scheduler_server/mod.rs +++ b/ballista/scheduler/src/scheduler_server/mod.rs @@ -31,7 +31,7 @@ use datafusion::prelude::{SessionConfig, SessionContext}; use datafusion_proto::logical_plan::AsLogicalPlan; use datafusion_proto::physical_plan::AsExecutionPlan; -use crate::cluster::{BallistaCluster, JobStateEventStream}; +use crate::cluster::{BallistaCluster, ClusterStateEventStream, JobStateEventStream}; use crate::config::SchedulerConfig; use crate::metrics::SchedulerMetricsCollector; use ballista_core::serde::scheduler::{ExecutorData, ExecutorMetadata}; @@ -188,6 +188,11 @@ impl SchedulerServer Result { + self.state.executor_manager.cluster_state_events().await + } + /// True when at least `min_ready_executors` executors currently have /// live heartbeats. Embedders can call this from their own health/readiness /// handler and AND it with whatever app-specific state they track. The @@ -456,6 +461,7 @@ pub fn timestamp_millis() -> u64 { #[cfg(test)] mod test { use std::sync::Arc; + use std::time::Duration; use ballista_core::JobId; use ballista_core::extension::SessionConfigExt; @@ -468,7 +474,9 @@ mod test { use datafusion::test_util::scan_empty_with_partitions; use datafusion_proto::protobuf::LogicalPlanNode; use datafusion_proto::protobuf::PhysicalPlanNode; + use futures::StreamExt; + use crate::cluster::ClusterStateEvent; use crate::config::SchedulerConfig; use ballista_core::config::TaskSchedulingPolicy; use ballista_core::error::Result; @@ -492,6 +500,31 @@ mod test { assert_submitted_event, test_cluster_context, }; + #[tokio::test] + async fn test_scheduler_exposes_cluster_state_events() -> Result<()> { + let scheduler = test_scheduler(TaskSchedulingPolicy::PushStaged).await?; + let mut events = scheduler.cluster_state_events().await?; + let (executor_metadata, executor_data) = + test_executors(2).into_iter().next().unwrap(); + let executor_id = executor_metadata.id.clone(); + + scheduler + .state + .executor_manager + .register_executor(executor_metadata, executor_data) + .await?; + + let event = tokio::time::timeout(Duration::from_secs(5), events.next()) + .await + .expect("cluster state event should arrive"); + assert_eq!( + event, + Some(ClusterStateEvent::RegisteredExecutor { executor_id }) + ); + + Ok(()) + } + #[tokio::test] async fn test_pull_scheduling() -> Result<()> { let plan = test_plan(); diff --git a/ballista/scheduler/src/state/executor_manager.rs b/ballista/scheduler/src/state/executor_manager.rs index 2f0b162b03..1da7565662 100644 --- a/ballista/scheduler/src/state/executor_manager.rs +++ b/ballista/scheduler/src/state/executor_manager.rs @@ -25,7 +25,7 @@ use ballista_core::serde::protobuf::ExecutorMetric; use ballista_core::serde::protobuf::executor_metric::Metric; use log::trace; -use crate::cluster::{BoundTask, ClusterState, ExecutorSlot}; +use crate::cluster::{BoundTask, ClusterState, ClusterStateEventStream, ExecutorSlot}; use crate::config::SchedulerConfig; use crate::state::execution_graph::RunningTaskInfo; @@ -112,6 +112,11 @@ impl ExecutorManager { Ok(()) } + /// Returns a stream of cluster state events from the configured state backend. + pub async fn cluster_state_events(&self) -> Result { + self.cluster_state.cluster_state_events().await + } + /// Binds ready-to-run tasks from active jobs to available executor slots. /// /// Returns a list of bound tasks that can be launched on executors.