diff --git a/ballista/core/src/ids.rs b/ballista/core/src/ids.rs index c10d23b37..1df1a6214 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 diff --git a/ballista/scheduler/src/config.rs b/ballista/scheduler/src/config.rs index 940888714..bbb7f5f6d 100644 --- a/ballista/scheduler/src/config.rs +++ b/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 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 533711ecb..e2259b5cf 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 000000000..612adab2c --- /dev/null +++ b/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`): 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 66ed01b50..b8435acab 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; @@ -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>, /// 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 { @@ -103,10 +111,12 @@ impl SchedulerServer SchedulerServer SchedulerServer SchedulerServer SchedulerServer broadcast::Receiver { + 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 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 ee5174853..c7ba08a13 100644 --- a/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs +++ b/ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs @@ -21,6 +21,7 @@ use std::time::Duration; use ballista_core::JobId; 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}; @@ -28,6 +29,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; @@ -45,6 +47,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 { @@ -52,14 +56,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); + } + async fn abort_job(&self, job_id: &JobId, failure_reason: String) -> Result<()> { let executor_manager = self.state.executor_manager.clone(); self.state @@ -124,7 +136,13 @@ 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 + let job_state_sender = self.job_state_sender.clone(); tokio::spawn(async move { let event = if let Err(e) = state .submit_job( @@ -171,6 +189,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, @@ -208,16 +228,20 @@ impl .record_failed(&job_id, queued_at, failed_at); error!("Job {job_id} failed: {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, @@ -228,9 +252,14 @@ impl .record_completed(&job_id, queued_at, completed_at); info!("Job finished successfully: [{job_id}]"); + let intermediate_stage_ids = match self.state.task_manager.succeed_job(&job_id).await { - Ok(ids) => ids, + Ok(ids) => { + // Broadcast only after state transition succeeds + self.broadcast_job_state(JobStateEvent::completed(&job_id)); + ids + } Err(e) => { error!( "Fail to invoke succeed_job for job {job_id} due to {e:?}" @@ -251,8 +280,17 @@ impl .record_failed(&job_id, queued_at, failed_at); error!("Job failed: [{job_id}]"); - if let Err(e) = self.abort_job(&job_id, fail_message).await { - error!("Fail to abort job {job_id} due to {e:?}"); + match self.abort_job(&job_id, fail_message.clone()).await { + Ok(()) => { + // Broadcast only after state transition succeeds + self.broadcast_job_state(JobStateEvent::failed( + &job_id, + &fail_message, + )); + } + Err(e) => { + error!("Fail to abort job {job_id} due to {e:?}"); + } } self.state.clean_up_failed_job(job_id); } @@ -266,8 +304,14 @@ impl self.metrics_collector.record_cancelled(&job_id); info!("Job cancelled: [{job_id}]"); - if let Err(e) = self.abort_job(&job_id, "Cancelled".to_owned()).await { - error!("Fail to cancel job {job_id} due to {e:?}"); + match self.abort_job(&job_id, "Cancelled".to_owned()).await { + Ok(()) => { + // Broadcast only after state transition succeeds + self.broadcast_job_state(JobStateEvent::cancelled(&job_id)); + } + Err(e) => { + error!("Fail to cancel job {job_id} due to {e:?}"); + } } self.state.clean_up_failed_job(job_id); } @@ -377,6 +421,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; @@ -441,6 +486,115 @@ 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 be184cd6e..2104c05e9 100644 --- a/ballista/scheduler/src/test_utils.rs +++ b/ballista/scheduler/src/test_utils.rs @@ -533,6 +533,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