1893: feat: add poll_now_notify to poll_loop and on_work_available callback - #77
1893: feat: add poll_now_notify to poll_loop and on_work_available callback#77martin-augment wants to merge 3 commits into
poll_now_notify to poll_loop and on_work_available callback#77Conversation
Adds an optional poll_now_notify: Option<Arc<Notify>> parameter to the executor poll_loop so the scheduler can wake an idle executor immediately instead of waiting for the next poll interval. The idle wait becomes a tokio::select! over the poll interval and the notify. Adds an OnWorkAvailableFn callback to SchedulerConfig, invoked from the query stage scheduler when new work becomes available (after JobSubmitted and when new stages become runnable). Ported from spiceai#12 for upstreaming.
WalkthroughThe PR updates multiple review-guidance files to forbid linking to GitHub issues or pull requests and to exclude AI-agent configuration files from review. It adds an optional ✨ Finishing Touches📝 Generate docstrings
🧪 Generate unit tests (beta)
Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out. Comment |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces a mechanism to notify executors immediately when new work becomes available, reducing polling latency. This is done by adding a poll_now_notify parameter to the executor's poll loop and an on_work_available callback to the scheduler. The review feedback points out three important improvements: refactoring an unstable let-chain condition to ensure stable Rust compatibility, adding documentation to warn against blocking operations in the synchronous callback, and increasing the idle poll interval when notifications are enabled to reduce gRPC load.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| if !stage_events.is_empty() | ||
| && let Some(ref callback) = self.config.on_work_available | ||
| { | ||
| callback("tasks_completed:new_stages_runnable"); | ||
| } |
There was a problem hiding this comment.
The use of let expressions in boolean conditions (let-chains) is an unstable Rust feature (let_chains) and will fail to compile on stable Rust.
To ensure compatibility with stable Rust, please refactor this to use nested if statements or a standard if let block.
| if !stage_events.is_empty() | |
| && let Some(ref callback) = self.config.on_work_available | |
| { | |
| callback("tasks_completed:new_stages_runnable"); | |
| } | |
| if !stage_events.is_empty() { | |
| if let Some(ref callback) = self.config.on_work_available { | |
| callback("tasks_completed:new_stages_runnable"); | |
| } | |
| } |
| /// Callback invoked when new work becomes available for executors. | ||
| /// | ||
| /// This is called after: | ||
| /// - A job is submitted and tasks are ready to be scheduled | ||
| /// - Tasks complete and new stages become runnable | ||
| /// | ||
| /// This allows external systems to notify executors to poll immediately | ||
| /// rather than waiting for their next poll interval. | ||
| pub type OnWorkAvailableFn = Arc<dyn Fn(&str) + Send + Sync>; |
There was a problem hiding this comment.
Since OnWorkAvailableFn is a synchronous callback invoked directly within the scheduler's main event loop (QueryStageScheduler::on_receive), any blocking operations or slow I/O performed inside this callback will block the entire event loop, severely degrading scheduler throughput and responsiveness.
Please add a documentation note warning users that implementations of this callback must be non-blocking and should delegate any heavy or blocking work to a separate thread or task (e.g., via tokio::spawn).
/// Callback invoked when new work becomes available for executors.
///
/// This is called after:
/// - A job is submitted and tasks are ready to be scheduled
/// - Tasks complete and new stages become runnable
///
/// This allows external systems to notify executors to poll immediately
/// rather than waiting for their next poll interval.
///
/// # Warning
///
/// This callback is executed synchronously within the scheduler's main event loop.
/// Implementations **must be non-blocking** and should offload any blocking or
/// long-running operations (such as network I/O) to a separate task or thread.
pub type OnWorkAvailableFn = Arc<dyn Fn(&str) + Send + Sync>;| if !active_job { | ||
| tokio::time::sleep(Duration::from_millis(50)).await; | ||
| // Wait for either the poll interval or a poll_now notification | ||
| match &poll_now_notify { | ||
| Some(notify) => { | ||
| tokio::select! { | ||
| () = tokio::time::sleep(Duration::from_millis(50)) => {} | ||
| () = notify.notified() => { | ||
| debug!("Received poll_now notification, polling immediately"); | ||
| } | ||
| } | ||
| } | ||
| None => { | ||
| tokio::time::sleep(Duration::from_millis(50)).await; | ||
| } | ||
| } | ||
| } |
There was a problem hiding this comment.
When poll_now_notify is provided, the executor can safely use a much longer idle poll interval (e.g., 1000ms or more) instead of the aggressive 50ms fallback.
Polling every 50ms from multiple executors can generate significant idle gRPC load on the scheduler. Since the Notify mechanism ensures near-instantaneous wakeups when work is actually available, increasing the sleep duration in the Some(notify) branch is a great opportunity to reduce idle overhead without sacrificing latency.
| if !active_job { | |
| tokio::time::sleep(Duration::from_millis(50)).await; | |
| // Wait for either the poll interval or a poll_now notification | |
| match &poll_now_notify { | |
| Some(notify) => { | |
| tokio::select! { | |
| () = tokio::time::sleep(Duration::from_millis(50)) => {} | |
| () = notify.notified() => { | |
| debug!("Received poll_now notification, polling immediately"); | |
| } | |
| } | |
| } | |
| None => { | |
| tokio::time::sleep(Duration::from_millis(50)).await; | |
| } | |
| } | |
| } | |
| if !active_job { | |
| // Wait for either the poll interval or a poll_now notification | |
| match &poll_now_notify { | |
| Some(notify) => { | |
| tokio::select! { | |
| () = tokio::time::sleep(Duration::from_millis(1000)) => {} | |
| () = notify.notified() => { | |
| debug!("Received poll_now notification, polling immediately"); | |
| } | |
| } | |
| } | |
| None => { | |
| tokio::time::sleep(Duration::from_millis(50)).await; | |
| } | |
| } | |
| } |
🤖 Augment PR SummarySummary: This PR adds a way to reduce executor polling latency by allowing explicit wake-ups. Changes:
Technical Notes: The callback is synchronous and runs on the scheduler event loop thread/task, so it should be kept lightweight. 🤖 Was this summary useful? React with 👍 or 👎 |
![augmentcode[bot]](https://avatars.githubusercontent.com/in/1027498?s=60&v=4){kind=small}
|
|
||
| // Notify external systems that new work is available | ||
| if let Some(ref callback) = self.config.on_work_available { | ||
| callback(&format!("job_submitted:{job_id}")); |
There was a problem hiding this comment.
on_work_available is user-provided code; if it blocks or panics here it can stall or crash the scheduler event loop. Consider isolating the callback so scheduler progress isn't coupled to external notification behavior.
Other locations where this applies: ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs:307
Severity: medium
Other Locations
ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs:307
🤖 Was this useful? React with 👍 or 👎, or 🚀 if it prevented an incident/outage.
![coderabbitai[bot]](https://avatars.githubusercontent.com/in/347564?s=60&v=4){kind=small}
There was a problem hiding this comment.
Actionable comments posted: 2
🤖 Prompt for all review comments with AI agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
Inline comments:
In `@ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs`:
- Around line 184-187: The `on_work_available` callback in
`QueryStageScheduler::on_receive` is being invoked inline on the scheduler event
loop, so move this notification work onto a separate task or channel instead of
calling it directly. Update the scheduler path to enqueue or spawn the callback
execution, and isolate any panic or failure so it cannot stall or crash
scheduling; apply the same pattern anywhere else
`SchedulerConfig::on_work_available` is called, including the other callback
site noted in the review.
- Around line 303-309: The wake-up callback in query_stage_scheduler.rs is
firing before the runnable-stage events are enqueued, so move the
on_work_available notification in QueryStageScheduler to after the stage_events
are pushed into the queue. Keep the existing stage_events check, but ensure the
enqueue path completes first and then invoke
callback("tasks_completed:new_stages_runnable") so idle executors can
immediately see the new runnable stages.
🪄 Autofix (Beta)
Fix all unresolved CodeRabbit comments on this PR:
- Push a commit to this branch (recommended)
- Create a new PR with the fixes
ℹ️ Review info
⚙️ Run configuration
Configuration used: Organization UI
Review profile: CHILL
Plan: Pro
Run ID: dd5fc870-0d93-47eb-8e71-a1210ff0c086
📒 Files selected for processing (10)
.cursor/rules.md.gemini/rules.mdAGENTS.mdCLAUDE.mdballista/executor/src/execution_loop.rsballista/executor/src/executor_process.rsballista/executor/src/standalone.rsballista/scheduler/src/config.rsballista/scheduler/src/scheduler_server/query_stage_scheduler.rsexamples/examples/mtls-cluster.rs
| // Notify external systems that new work is available | ||
| if let Some(ref callback) = self.config.on_work_available { | ||
| callback(&format!("job_submitted:{job_id}")); | ||
| } |
There was a problem hiding this comment.
🩺 Stability & Availability | 🟠 Major | 🏗️ Heavy lift
Don’t run user-provided callbacks inline on the scheduler event loop.
on_work_available comes from SchedulerConfig, so this is arbitrary external code. Calling it directly in on_receive means a slow notifier stalls scheduling, and a panic in the callback can take down the event-loop task. Hand the notification off to a separate task/channel and isolate failures from the scheduler path.
Also applies to: 303-309
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs` around
lines 184 - 187, The `on_work_available` callback in
`QueryStageScheduler::on_receive` is being invoked inline on the scheduler event
loop, so move this notification work onto a separate task or channel instead of
calling it directly. Update the scheduler path to enqueue or spawn the callback
execution, and isolate any panic or failure so it cannot stall or crash
scheduling; apply the same pattern anywhere else
`SchedulerConfig::on_work_available` is called, including the other callback
site noted in the review.
| // Notify external systems when new stages become runnable | ||
| if !stage_events.is_empty() | ||
| && let Some(ref callback) = self.config.on_work_available | ||
| { | ||
| callback("tasks_completed:new_stages_runnable"); | ||
| } | ||
|
|
There was a problem hiding this comment.
🎯 Functional Correctness | 🟠 Major | ⚡ Quick win
Notify only after enqueuing the runnable-stage events.
Line 303 fires the wake-up before Lines 310-312 enqueue the stage_events. That lets an idle executor poll before the new runnable stages exist, which defeats the “poll now” fast path and can fall back to the next interval anyway.
Suggested fix
- // Notify external systems when new stages become runnable
- if !stage_events.is_empty()
- && let Some(ref callback) = self.config.on_work_available
- {
- callback("tasks_completed:new_stages_runnable");
- }
-
- for stage_event in stage_events {
+ let has_new_runnable_stages = !stage_events.is_empty();
+ for stage_event in stage_events {
event_sender.post_event(stage_event).await?;
}
+ if has_new_runnable_stages
+ && let Some(ref callback) = self.config.on_work_available
+ {
+ callback("tasks_completed:new_stages_runnable");
+ }📝 Committable suggestion
‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.
| // Notify external systems when new stages become runnable | |
| if !stage_events.is_empty() | |
| && let Some(ref callback) = self.config.on_work_available | |
| { | |
| callback("tasks_completed:new_stages_runnable"); | |
| } | |
| let has_new_runnable_stages = !stage_events.is_empty(); | |
| for stage_event in stage_events { | |
| event_sender.post_event(stage_event).await?; | |
| } | |
| if has_new_runnable_stages | |
| && let Some(ref callback) = self.config.on_work_available | |
| { | |
| callback("tasks_completed:new_stages_runnable"); | |
| } |
🤖 Prompt for AI Agents
Verify each finding against current code. Fix only still-valid issues, skip the
rest with a brief reason, keep changes minimal, and validate.
In `@ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs` around
lines 303 - 309, The wake-up callback in query_stage_scheduler.rs is firing
before the runnable-stage events are enqueued, so move the on_work_available
notification in QueryStageScheduler to after the stage_events are pushed into
the queue. Keep the existing stage_events check, but ensure the enqueue path
completes first and then invoke callback("tasks_completed:new_stages_runnable")
so idle executors can immediately see the new runnable stages.
Code ReviewThis PR adds two building blocks for reducing executor poll latency: a
|
3 participants
1893: To review by AI