1893: feat: add poll_now_notify to poll_loop and on_work_available callback

.cursor/rules.md

@@ -0,0 +1,5 @@
+When making a code review never link to any issue or a pull request
+in any GitHub repository!
+
+Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md!
+

.gemini/rules.md

@@ -0,0 +1,5 @@
+When making a code review never link to any issue or a pull request
+in any GitHub repository!
+
+Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md!
+

AGENTS.md

@@ -0,0 +1,5 @@
+When making a code review never link to any issue or a pull request
+in any GitHub repository!
+
+Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md!
+

CLAUDE.md

@@ -0,0 +1,5 @@
+When making a code review never link to any issue or a pull request
+in any GitHub repository!
+
+Do not review AI agents' config files like CLAUDE.md, AGENTS.md or rules.md!
+

ballista/executor/src/execution_loop.rs

@@ -46,7 +46,7 @@ use std::error::Error;
 use std::sync::mpsc::{Receiver, Sender, TryRecvError};
 use std::time::{SystemTime, UNIX_EPOCH};
 use std::{sync::Arc, time::Duration};
-use tokio::sync::{OwnedSemaphorePermit, Semaphore};
+use tokio::sync::{Notify, OwnedSemaphorePermit, Semaphore};
 use tonic::codegen::{Body, Bytes, StdError};
 
 /// Main execution loop that polls the scheduler for available tasks.
@@ -57,10 +57,14 @@ use tonic::codegen::{Body, Bytes, StdError};
 ///
 /// The loop respects the executor's concurrent task limit via a semaphore,
 /// ensuring no more than the configured number of tasks run simultaneously.
+///
+/// `poll_now_notify`, when provided, allows the scheduler to wake the poll loop
+/// immediately instead of waiting for the next idle poll interval.
 pub async fn poll_loop<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan, C>(
     mut scheduler: SchedulerGrpcClient<C>,
     executor: Arc<Executor>,
     codec: BallistaCodec<T, U>,
+    poll_now_notify: Option<Arc<Notify>>,
 ) -> Result<(), BallistaError>
 where
     C: tonic::client::GrpcService<tonic::body::Body>,
@@ -204,7 +208,20 @@ where
         }
 
         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;
+                }
+            }
         }
     }
 }

ballista/executor/src/executor_process.rs

@@ -505,6 +505,7 @@ pub async fn start_executor_process(
                 scheduler.clone(),
                 executor.clone(),
                 default_codec,
+                None, // poll_now_notify
             )));
         }
     };

ballista/executor/src/standalone.rs

@@ -142,7 +142,7 @@ pub async fn new_standalone_executor_from_builder(
             )),
     );
 
-    tokio::spawn(execution_loop::poll_loop(scheduler, executor, codec));
+    tokio::spawn(execution_loop::poll_loop(scheduler, executor, codec, None));
     Ok(())
 }
 

ballista/scheduler/src/config.rs

@@ -34,6 +34,16 @@ use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use std::fmt::Display;
 use std::sync::Arc;
 
+/// 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>;
+
 /// Command-line configuration for the scheduler binary.
 #[cfg(feature = "build-binary")]
 #[derive(clap::Parser, Debug)]
@@ -278,6 +288,9 @@ pub struct SchedulerConfig {
     #[cfg(feature = "rest-api")]
     /// Comma-separated list of allowed methods for CORS
     pub cors_allowed_methods: String,
+    /// Callback invoked when new work becomes available for executors.
+    /// The string argument is a reason/description for debugging purposes.
+    pub on_work_available: Option<OnWorkAvailableFn>,
 }
 
 impl Default for SchedulerConfig {
@@ -314,6 +327,7 @@ impl Default for SchedulerConfig {
             cors_allowed_origins: String::default(),
             #[cfg(feature = "rest-api")]
             cors_allowed_methods: String::default(),
+            on_work_available: None,
         }
     }
 }
@@ -546,6 +560,7 @@ impl TryFrom<Config> for SchedulerConfig {
             cors_allowed_origins: opt.cors_allowed_origins,
             #[cfg(feature = "rest-api")]
             cors_allowed_methods: opt.cors_allowed_methods,
+            on_work_available: None,
         };
 
         Ok(config)

ballista/scheduler/src/scheduler_server/query_stage_scheduler.rs

@@ -180,6 +180,11 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
                         .post_event(QueryStageSchedulerEvent::ReviveOffers)
                         .await?;
                 }
+
+                // Notify external systems that new work is available
+                if let Some(ref callback) = self.config.on_work_available {
+                    callback(&format!("job_submitted:{job_id}"));
+                }
             }
             QueryStageSchedulerEvent::JobPlanningFailed {
                 job_id,
@@ -295,6 +300,13 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
                                 .await?;
                         }
 
+                        // 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 {
                             event_sender.post_event(stage_event).await?;
                         }

examples/examples/mtls-cluster.rs

@@ -430,7 +430,7 @@ async fn run_executor() -> Result<(), Box<dyn std::error::Error>> {
     // This registers the executor and starts polling for tasks
     info!("Starting execution poll loop...");
     let poll_handle = tokio::spawn(async move {
-        execution_loop::poll_loop(scheduler, executor, codec).await
+        execution_loop::poll_loop(scheduler, executor, codec, None).await
     });
 
     tokio::select! {