feat: add `poll_now_notify` to `poll_loop` and `on_work_available` callback by Jeadie · Pull Request #1893 · apache/datafusion-ballista

Jeadie · 2026-06-23T05:43:45Z

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?

Adds two related hooks for push-style executor scheduling:

An optional poll_now_notify: Option<Arc<Notify>> parameter on the executor poll_loop. When provided, the idle wait becomes a tokio::select! over the poll interval and the notify, so the scheduler can wake an idle executor immediately instead of waiting for the next poll tick. None preserves the existing timer-only behavior.
An OnWorkAvailableFn callback on SchedulerConfig (on_work_available), invoked by the query stage scheduler when new work becomes available (after JobSubmitted and when new stages become runnable).

Ported from spiceai#12

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 #12 for upstreaming.

…andalone

martin-g · 2026-06-26T08:58:48Z

+/// - 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.


Suggested change

/// rather than waiting for their next poll interval.

/// 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.

The callback is synchronous and a blocking implementation may affect badly the Scheduler;w work.

Comment added.

martin-g · 2026-06-26T09:03:00Z

+            match &poll_now_notify {
+                Some(notify) => {
+                    tokio::select! {
+                        () = tokio::time::sleep(Duration::from_millis(50)) => {}


The duration here could be longer.
The idea is that notify will notify us as soon as there is some work, right ?
So, there is no need to poll so frequently if an external mechanism will notify us when to do it.

Agreed, implemented to use a 1s fallback if a poll_now_notify is provided, otherwise unchanged at 50ms to preserve existing behavior.

martin-g · 2026-06-26T09:10:45Z

+                        // 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");
+                        }
+


Shouldn't this be done after line 312 ? I.e. after posting the events.
Currently this will tell executors to poll but it seems to early.

The stage_events are not necessarily "new stages are runnable". They might be JobRunningFailed or JobFinished, neither of which represents new runnable work.

This is correct, changed to remove the callback from TaskUpdating and moved into the JobUpdated arm, assuming update_job returns new_tasks > 0. So this only notifies when stages genuinely had new scheduleable tasks, and only after they are updated.

martin-g · 2026-06-26T09:13:22Z

            cors_allowed_origins: opt.cors_allowed_origins,
            #[cfg(feature = "rest-api")]
            cors_allowed_methods: opt.cors_allowed_methods,
+            on_work_available: None,


How is this new setting supposed to be set ?
I'd expect at least a new with_on_work_available setter to be added if there is no new command line option.

Added SchedulerConfig::with_on_work_available

martin-g · 2026-06-26T09:17:52Z

+
+                // Notify external systems that new work is available
+                if let Some(ref callback) = self.config.on_work_available {
+                    callback(&format!("job_submitted:{job_id}"));


The event_sender.post_event(QueryStageSchedulerEvent::ReviveOffers) is non-blocking and the events might be still in the channel when this callback is executed. I.e. the first poll may see no events.

Fixed. The callback now fires only after graph.revive() has run (in update_job, and submit_plan revives before caching), which is the exact point tasks become visible to poll_work. Now a woken executor's first poll always finds the work.

martin-g · 2026-06-26T09:21:12Z

+///
+/// 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>;


The function receives a plain &str as a command/message. IMO it would be better to use an Enum is all possible variants of commands.

Is there a need to use an Arc ?
It could be a Box.

Replaced the &str with a WorkAvailableReason enum (JobSubmitted { job_id } / NewStagesRunnable { job_id }), so consumers get structured data instead of parsing strings.

It also needs to stay Arc because SchedulerConfig derives Clone and Box<dyn Fn> isn't cloneable. We could drop the Clone requirement if we want to keep it as Box.

- Replace the callback's `&str` argument with a `WorkAvailableReason` enum - Fire the stage-resolution callback from the JobUpdated arm after `update_job` has revived the graph, so it only fires for genuinely new runnable work and never before the work is pollable - Add `SchedulerConfig::with_on_work_available` - Document that the callback runs synchronously in the scheduler event loop and must not block - Idle-poll less aggressively (1s fallback) when a poll_now notifier is wired

phillipleblanc · 2026-07-20T01:29:26Z

@martin-g Thanks for the thorough review, this should be ready now for another look.

Jeadie added 2 commits June 23, 2026 13:58

fix(examples): pass new poll_loop arg in mtls-cluster and reformat st…

683dfa7

…andalone

Jeadie mentioned this pull request Jun 23, 2026

feat: add poll_now_notify to poll_loop and on_work_available callback spiceai/datafusion-ballista#50

Open

Jeadie changed the title ~~Upstream/poll now notify on work available~~ feat: add poll_now_notify to poll_loop and on_work_available callback Jun 23, 2026

martin-g reviewed Jun 26, 2026

View reviewed changes

phillipleblanc added 2 commits July 20, 2026 10:09

Merge branch 'main' into upstream/poll-now-notify-on-work-available

e7b8ed3

-/// rather than waiting for their next poll interval.
+/// 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.

Uh oh!

Conversation

Jeadie commented Jun 23, 2026

Which issue does this PR close?

Rationale for this change

What changes are included in this PR?

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

phillipleblanc commented Jul 20, 2026

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants