feat(scheduler): expose existing job state event stream

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/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,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<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;
@@ -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<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> {
+    /// 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<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             scheduler_name.clone(),
             config.clone(),
         ));
+        let (job_state_sender, _) = broadcast::channel(Self::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(),
@@ -121,6 +137,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             #[cfg(feature = "rest-api")]
             query_stage_scheduler,
             config,
+            job_state_sender,
         }
     }
 
@@ -141,10 +158,12 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             config.clone(),
             task_launcher,
         ));
+        let (job_state_sender, _) = broadcast::channel(Self::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(),
@@ -160,6 +179,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             #[cfg(feature = "rest-api")]
             query_stage_scheduler,
             config,
+            job_state_sender,
         }
     }
 
@@ -181,6 +201,37 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
     pub fn running_job_number(&self) -> 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<job_state_event::JobStateEvent> {
+        self.job_state_sender.subscribe()
+    }
+
     #[cfg(feature = "rest-api")]
     pub(crate) fn metrics_collector(&self) -> &dyn SchedulerMetricsCollector {
         self.query_stage_scheduler.metrics_collector()