feat: add a broadcast channel for job state event notifications#1891
feat: add a broadcast channel for job state event notifications#1891Jeadie wants to merge 5 commits into
Conversation
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. Ported from #15 for upstreaming.
andygrove
left a comment
andygrove
left a comment
There was a problem hiding this comment.
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.
| /// | ||
| /// This determines how many job state events can be buffered before | ||
| /// slow receivers start lagging behind. | ||
| const JOB_STATE_CHANNEL_CAPACITY: usize = 256; |
There was a problem hiding this comment.
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.
There was a problem hiding this comment.
This is now configurable.
…nnel - Add job_state_channel_capacity to SchedulerConfig (default 256) so operators with high job throughput can tune the buffer; removes the hardcoded const from SchedulerServer - Move all job state broadcasts to after their state transitions succeed, so subscribers never see an event for a commit that failed; aligns Queued/Failed/Completed/Cancelled with how Running already behaved - Add From<job_status::Status> for JobState to keep the two enums in sync; document that Cancelled has no protobuf counterpart - Expand module doc with guidance on when to use JobStateEvent vs JobStatusSubscriber, and reference job_state_channel_capacity for tuning - Add integration tests covering the broadcast wiring: lifecycle order (Queued -> Running -> Completed) and cancel path
|
@andygrove this is ready for another review, thanks! |
3 participants
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::broadcastsender toSchedulerServer, subscribable viasubscribe_job_updates(). On job state changes (queued / running / completed / failed / cancelled) the scheduler broadcasts aJobStateEventcarrying 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 throughsubmit_job(..., subscriber). This broadcast channel is a separate, fan-out style notification path; worth deciding the best DX between these.JobStatusSubscriber(existing)JobStateEventbroadcast (this PR)mpsc- single consumerbroadcast- fan-outExecuteQuery/GetJobStatusResult.subscribe_job_updates()JobStatusJobStateEvent { job_id: String, state }Both fire from the same lifecycle points in
query_stage_scheduler.rs.Ported from spiceai#15.