feat(scheduler): expose existing job state event stream

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(

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;
+use crate::cluster::{BallistaCluster, JobStateEventStream};
 use crate::config::SchedulerConfig;
 use crate::metrics::SchedulerMetricsCollector;
 use ballista_core::serde::scheduler::{ExecutorData, ExecutorMetadata};
@@ -183,6 +183,11 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
         self.state.task_manager.running_job_number()
     }
 
+    /// Returns a stream of job state events from the configured state backend.
+    pub async fn job_state_events(&self) -> Result<JobStateEventStream> {
+        self.state.task_manager.job_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

ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs

@@ -465,6 +465,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
 
 #[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);

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<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> TaskManager<T, U>
             .count()
     }
 
+    /// Returns a stream of job state events from the configured state backend.
+    pub async fn job_state_events(&self) -> Result<JobStateEventStream> {
+        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

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<JobStateEventStream> {
+        self.scheduler.job_state_events().await
+    }
+
     /// Returns the session context for tests.
     pub async fn ctx(&self) -> Result<Arc<SessionContext>> {
         self.scheduler