feat: add a broadcast channel for job state event notifications

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<str> for $name {
             fn as_ref(&self) -> &str {
                 &self.0

ballista/scheduler/src/config.rs

@@ -278,6 +278,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,
@@ -318,6 +323,7 @@ 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")]
@@ -459,6 +465,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
@@ -551,6 +563,7 @@ impl TryFrom<Config> 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")]

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};

ballista/scheduler/src/scheduler_server/job_state_event.rs

@@ -0,0 +1,186 @@
+// 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<JobStatus>`): 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<job_status::Status>` 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<job_status::Status> 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<String>, 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<String>) -> Self {
+        Self::new(job_id, JobState::Queued)
+    }
+
+    /// Creates a running event for the given job.
+    pub fn running(job_id: impl Into<String>) -> Self {
+        Self::new(job_id, JobState::Running)
+    }
+
+    /// Creates a completed event for the given job.
+    pub fn completed(job_id: impl Into<String>) -> Self {
+        Self::new(job_id, JobState::Completed)
+    }
+
+    /// Creates a failed event for the given job.
+    pub fn failed(job_id: impl Into<String>, error: impl Into<String>) -> Self {
+        Self::new(job_id, JobState::Failed(error.into()))
+    }
+
+    /// Creates a cancelled event for the given job.
+    pub fn cancelled(job_id: impl Into<String>) -> 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");
+    }
+}

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;
@@ -58,6 +59,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.
@@ -86,6 +89,11 @@ pub struct SchedulerServer<T: 'static + AsLogicalPlan, U: 'static + AsExecutionP
     query_stage_scheduler: Arc<QueryStageScheduler<T, U>>,
     /// Scheduler configuration.
     config: Arc<SchedulerConfig>,
+    /// 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<job_state_event::JobStateEvent>,
 }
 
 impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T, U> {
@@ -103,10 +111,12 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             scheduler_name.clone(),
             config.clone(),
         ));
+        let (job_state_sender, _) = broadcast::channel(config.job_state_channel_capacity);
         let query_stage_scheduler = Arc::new(QueryStageScheduler::new(
             state.clone(),
             metrics_collector,
             config.clone(),
+            job_state_sender.clone(),
         ));
         let query_stage_event_loop = EventLoop::new(
             "query_stage".to_owned(),
@@ -122,6 +132,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             #[cfg(feature = "rest-api")]
             query_stage_scheduler,
             config,
+            job_state_sender,
         }
     }
 
@@ -142,10 +153,12 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             config.clone(),
             task_launcher,
         ));
+        let (job_state_sender, _) = broadcast::channel(config.job_state_channel_capacity);
         let query_stage_scheduler = Arc::new(QueryStageScheduler::new(
             state.clone(),
             metrics_collector,
             config.clone(),
+            job_state_sender.clone(),
         ));
         let query_stage_event_loop = EventLoop::new(
             "query_stage".to_owned(),
@@ -161,6 +174,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             #[cfg(feature = "rest-api")]
             query_stage_scheduler,
             config,
+            job_state_sender,
         }
     }
 
@@ -183,6 +197,36 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
         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<job_state_event::JobStateEvent> {
+        self.job_state_sender.subscribe()
+    }
+
     /// 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
@@ -191,6 +235,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
         let alive = self.state.executor_manager.get_alive_executors().len();
         alive >= self.state.config.min_ready_executors
     }
+
     #[cfg(feature = "rest-api")]
     pub(crate) fn metrics_collector(&self) -> &dyn SchedulerMetricsCollector {
         self.query_stage_scheduler.metrics_collector()