feat(scheduler): expose existing job state event stream

[Jeadie]

Which issue does this PR close?

None Closes #. Suggested by @milenkovicm that we should upstream our changes (see comment).

Rationale for this change

Useful in spiceai/datafusion-ballista for production features.

What changes are included in this PR?

Summary

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. Lets components holding a handle to the scheduler observe job lifecycle transitions without polling.

Relationship to the existing JobStatusSubscriber

Similar to the existing JobStatusSubscriber (mpsc) mechanism wired through submit_job(..., subscriber). This broadcast channel is a separate, fan-out style notification path; worth deciding the best DX between these.

	JobStatusSubscriber (existing)	JobStateEvent broadcast (this PR)
Channel	mpsc - single consumer	broadcast - fan-out
Scope	1 subscriber per submitted job, request-scoped	Global; fires for all jobs
Consumer	submitting gRPC client, via ExecuteQuery / GetJobStatusResult.	in-process component(s), via subscribe_job_updates()
Payload	Full protobuf JobStatus	Lightweight JobStateEvent { job_id: String, state }

Both fire from the same lifecycle points in query_stage_scheduler.rs.

Ported from spiceai#15.

[andygrove]

Thanks for upstreaming this. It's cleanly structured, the doc comments are thorough, and I like that it's purely additive with no change to existing behavior. Ignoring send errors when there are no receivers is the right call for a broadcast channel.

On the design question you raised yourself about this versus the existing JobStatusSubscriber: agreed it's worth settling. It might help to document when a consumer should reach for each one, since they overlap in purpose but differ in scope and payload. If the broadcast path can eventually cover the in-process cases the mpsc path serves, consolidating later would be nice, but that doesn't need to block this.

Would it also be worth a test that subscribes via subscribe_job_updates() and drives a job through its lifecycle, asserting the expected events arrive? The unit tests cover the event types well, and a test on the actual broadcast wiring would guard the lifecycle hookups in query_stage_scheduler.rs against future refactors.

One minor note: JobState mirrors the protobuf job_status::Status. A From conversion between them, or a comment tying them together, could help keep the two from drifting over time.

This review was prepared with the help of an LLM.

[andygrove]

The 256 capacity is a sensible default. For a scheduler with high job throughput and a slow or briefly stalled subscriber, the buffer could fill and subscribers would start seeing Lagged. Would it be worth pulling this into SchedulerConfig so operators with heavier job rates can size it, keeping 256 as the default? Different deployments run very different workloads, so a fixed buffer may not fit everyone.

[phillipleblanc]

This is now configurable.

[phillipleblanc]

@andygrove this is ready for another review, thanks!

[milenkovicm]

apologies @phillipleblanc for slow review, i'll try to catch up
one question, current ballista implementation does have

JobState::job_state_events(&self) -> Result<JobStateEventStream>;

not its a bit of ceremony to subscribe to it.

would it make sense for TaskManager to expose

    pub async fn job_state_events(&self) -> Result<JobStateEventStream> {
        self.state.job_state_events().await
    }

and then SchedulerServer to expose

pub async fn job_state_events(&self) -> Result<JobStateEventStream> {
        self.state.task_manager.job_state_events().await
    }

JobStateEventStream is defined as:

pub type JobStateEventStream = Pin<Box<dyn Stream<Item = JobStateEvent> + Send>>; 

[milenkovicm]

can we align on one, please. i believe @ akshaychitneni is using existing one

[phillipleblanc]

@milenkovicm Agreed, I've pushed a change that exposes the existing stream rather than adding a new one.

[milenkovicm]

thanks @phillipleblanc

can i ask if you could expose another method, ClusterState has cluster_state_events

    async fn cluster_state_events(&self) -> Result<ClusterStateEventStream> {
        Ok(Box::pin(self.cluster_event_sender.subscribe()))
    }

which emits cluster related notifications, could you please expose that one to Scheduler as well

[milenkovicm]

or if you prefer we can do it in another PR, whatever suits @phillipleblanc

[phillipleblanc]

I'll do it in this PR, I need to merge main anyway

[milenkovicm]

thanks @phillipleblanc