feat: add poll_now_notify to poll_loop and on_work_available callback

ballista/executor/src/execution_loop.rs

@@ -46,9 +46,16 @@ use std::error::Error;
 use std::sync::mpsc::{Receiver, Sender, TryRecvError};
 use std::time::{Instant, 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};
 
+/// Idle sleep between polls when polling is the only way to learn of new work.
+const IDLE_POLL_INTERVAL: Duration = Duration::from_millis(50);
+
+/// Idle sleep when a `poll_now_notify` wake-up is wired and the timer is only
+/// a fallback.
+const NOTIFIED_IDLE_POLL_INTERVAL: Duration = Duration::from_secs(1);
+
 /// Maximum time the poll loop waits for a free vcore before polling the
 /// scheduler anyway. `poll_work` doubles as the executor's heartbeat under
 /// pull-based scheduling, so a fully-busy executor must keep polling (reporting
@@ -68,6 +75,12 @@ const HEARTBEAT_POLL_INTERVAL: Duration = Duration::from_secs(5);
 /// Pass `None` to have the loop create a semaphore sized to the executor's
 /// configured vcore count.
 ///
+/// `poll_now_notify`, when provided, wakes an idle poll loop immediately
+/// (typically wired to the scheduler's `on_work_available` callback) instead
+/// of waiting out the idle interval. A notification sent mid-poll is not
+/// lost: `Notify` stores the permit and the next `notified().await` returns
+/// immediately.
+///
 /// **Shared semaphores**: when one semaphore is shared across loops that
 /// connect to different schedulers, each scheduler independently sees the
 /// current free capacity and may dispatch up to that many tasks. The semaphore
@@ -89,6 +102,7 @@ pub async fn poll_loop<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan,
     mut scheduler: SchedulerGrpcClient<C>,
     executor: Arc<Executor>,
     codec: BallistaCodec<T, U>,
+    poll_now_notify: Option<Arc<Notify>>,
     free_vcores: Option<Arc<Semaphore>>,
     health: crate::health::ExecutorHealth,
 ) -> Result<(), BallistaError>
@@ -263,7 +277,19 @@ where
         }
 
         if !active_job {
-            tokio::time::sleep(Duration::from_millis(50)).await;
+            match &poll_now_notify {
+                Some(notify) => {
+                    tokio::select! {
+                        () = tokio::time::sleep(NOTIFIED_IDLE_POLL_INTERVAL) => {}
+                        () = notify.notified() => {
+                            debug!("Received poll_now notification, polling immediately");
+                        }
+                    }
+                }
+                None => {
+                    tokio::time::sleep(IDLE_POLL_INTERVAL).await;
+                }
+            }
         }
     }
 }

ballista/executor/src/executor_process.rs

@@ -728,6 +728,7 @@ pub async fn start_executor_process(
                 scheduler.clone(),
                 executor.clone(),
                 default_codec,
+                None, // poll_now_notify
                 None, // free_vcores: use internal semaphore
                 health,
             )));

ballista/executor/src/standalone.rs

@@ -148,6 +148,7 @@ pub async fn new_standalone_executor_from_builder(
         executor,
         codec,
         None,
+        None,
         crate::health::ExecutorHealth::new(),
     ));
     Ok(())

ballista/scheduler/src/config.rs

@@ -28,12 +28,43 @@
 use crate::SessionBuilder;
 use crate::cluster::DistributionPolicy;
 use ballista_core::extension::EndpointOverrideFn;
-use ballista_core::{ConfigProducer, config::TaskSchedulingPolicy};
+use ballista_core::{ConfigProducer, JobId, config::TaskSchedulingPolicy};
 use datafusion_proto::logical_plan::LogicalExtensionCodec;
 use datafusion_proto::physical_plan::PhysicalExtensionCodec;
 use std::fmt::Display;
 use std::sync::Arc;
 
+/// Why the scheduler believes new work has become available for executors.
+#[derive(Debug, Clone, PartialEq, Eq)]
+pub enum WorkAvailableReason {
+    /// A job was submitted and its initial tasks are ready to be scheduled.
+    JobSubmitted {
+        /// Identifier of the submitted job.
+        job_id: JobId,
+    },
+    /// Completed tasks resolved downstream stages of a job, and the tasks of
+    /// those stages are now schedulable.
+    NewStagesRunnable {
+        /// Identifier of the job that gained schedulable tasks.
+        job_id: JobId,
+    },
+}
+
+/// Callback invoked when new work becomes available for executors, e.g. to
+/// wake idle pull-based executors via the poll loop's `poll_now_notify`.
+///
+/// It fires only after the work is visible to a polling executor, so waking
+/// one cannot race the scheduler's internal event processing.
+///
+/// # Warning
+///
+/// The callback runs synchronously inside the scheduler's main event loop.
+/// Implementations **must be non-blocking**; offload blocking or long-running
+/// work (such as network I/O) to a separate task or thread.
+///
+/// `Arc` rather than `Box` so [`SchedulerConfig`] remains [`Clone`].
+pub type OnWorkAvailableFn = Arc<dyn Fn(WorkAvailableReason) + Send + Sync>;
+
 /// Command-line configuration for the scheduler binary.
 #[cfg(feature = "build-binary")]
 #[derive(clap::Parser, Debug)]
@@ -313,6 +344,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.
+    /// See [`OnWorkAvailableFn`].
+    pub on_work_available: Option<OnWorkAvailableFn>,
 }
 
 impl Default for SchedulerConfig {
@@ -352,6 +386,7 @@ impl Default for SchedulerConfig {
             cors_allowed_origins: String::default(),
             #[cfg(feature = "rest-api")]
             cors_allowed_methods: String::default(),
+            on_work_available: None,
         }
     }
 }
@@ -500,6 +535,16 @@ impl SchedulerConfig {
         self.use_tls = use_tls;
         self
     }
+
+    /// Sets the callback invoked when new work becomes available for
+    /// executors. See [`OnWorkAvailableFn`].
+    pub fn with_on_work_available(
+        mut self,
+        on_work_available: OnWorkAvailableFn,
+    ) -> Self {
+        self.on_work_available = Some(on_work_available);
+        self
+    }
 }
 
 /// Policy of distributing tasks to available executor slots
@@ -600,6 +645,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

@@ -27,7 +27,7 @@ use ballista_core::error::{BallistaError, Result};
 use ballista_core::event_loop::{EventAction, EventSender};
 use tokio::sync::mpsc::error::TrySendError;
 
-use crate::config::SchedulerConfig;
+use crate::config::{SchedulerConfig, WorkAvailableReason};
 use crate::metrics::SchedulerMetricsCollector;
 use crate::scheduler_server::timestamp_millis;
 use datafusion_proto::logical_plan::AsLogicalPlan;
@@ -204,6 +204,12 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
                         .post_event(QueryStageSchedulerEvent::ReviveOffers)
                         .await?;
                 }
+
+                // The graph was revived before caching, so the job's tasks
+                // are already visible to polling executors.
+                if let Some(callback) = &self.config.on_work_available {
+                    callback(WorkAvailableReason::JobSubmitted { job_id });
+                }
             }
             QueryStageSchedulerEvent::JobPlanningFailed {
                 job_id,
@@ -265,8 +271,19 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
             }
             QueryStageSchedulerEvent::JobUpdated(job_id) => {
                 debug!("Job updated, job_id: [{job_id}]");
-                if let Err(e) = self.state.task_manager.update_job(&job_id).await {
-                    error!("Fail to invoke update_job for job {job_id} due to {e:?}");
+                match self.state.task_manager.update_job(&job_id).await {
+                    Ok(new_tasks) => {
+                        // update_job revived the graph: the new tasks are
+                        // already visible to polling executors.
+                        if new_tasks > 0
+                            && let Some(callback) = &self.config.on_work_available
+                        {
+                            callback(WorkAvailableReason::NewStagesRunnable { job_id });
+                        }
+                    }
+                    Err(e) => {
+                        error!("Fail to invoke update_job for job {job_id} due to {e:?}");
+                    }
                 }
             }
             QueryStageSchedulerEvent::JobCancel(job_id) => {
@@ -465,17 +482,34 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
 
 #[cfg(test)]
 mod tests {
-    use crate::config::SchedulerConfig;
-    use crate::test_utils::{SchedulerTest, TestMetricsCollector, await_condition};
+    use crate::config::{SchedulerConfig, WorkAvailableReason};
+    use crate::scheduler_server::SchedulerServer;
+    use crate::test_utils::{
+        SchedulerTest, TestMetricsCollector, await_condition, test_cluster_context,
+    };
+    use ballista_core::BALLISTA_PROTOCOL_VERSION;
     use ballista_core::config::TaskSchedulingPolicy;
     use ballista_core::error::Result;
+    use ballista_core::extension::SessionConfigExt;
+    use ballista_core::serde::BallistaCodec;
     use ballista_core::serde::protobuf::job_status;
+    use ballista_core::serde::protobuf::scheduler_grpc_server::SchedulerGrpc;
+    use ballista_core::serde::protobuf::{
+        ExecutorRegistration, PollWorkParams, ShuffleWritePartition, SuccessfulTask,
+        TaskStatus, task_status,
+    };
+    use ballista_core::serde::scheduler::{
+        ExecutorOperatingSystemSpecification, ExecutorSpecification,
+    };
     use datafusion::arrow::datatypes::{DataType, Field, Schema};
     use datafusion::functions_aggregate::sum::sum;
     use datafusion::logical_expr::{LogicalPlan, col};
+    use datafusion::prelude::SessionConfig;
     use datafusion::test_util::scan_empty_with_partitions;
-    use std::sync::Arc;
+    use datafusion_proto::protobuf::{LogicalPlanNode, PhysicalPlanNode};
+    use std::sync::{Arc, Mutex};
     use std::time::Duration;
+    use tonic::Request;
 
     #[tokio::test]
     async fn test_pending_job_metric() -> Result<()> {
@@ -531,6 +565,136 @@ mod tests {
         Ok(())
     }
 
+    #[tokio::test]
+    async fn test_on_work_available_callback() -> Result<()> {
+        let reasons: Arc<Mutex<Vec<WorkAvailableReason>>> = Arc::default();
+        let captured = reasons.clone();
+
+        let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
+            SchedulerServer::new(
+                "localhost:50050".to_owned(),
+                test_cluster_context(),
+                BallistaCodec::default(),
+                Arc::new(
+                    SchedulerConfig {
+                        scheduling_policy: TaskSchedulingPolicy::PullStaged,
+                        ..Default::default()
+                    }
+                    .with_on_work_available(Arc::new(
+                        move |reason| {
+                            captured.lock().unwrap().push(reason);
+                        },
+                    )),
+                ),
+                Arc::new(TestMetricsCollector::default()),
+            );
+        scheduler.init().await?;
+
+        let ctx = scheduler
+            .state
+            .session_manager
+            .create_or_update_session(
+                "session",
+                &SessionConfig::new_with_ballista().with_target_partitions(2),
+            )
+            .await?;
+        let job_id = scheduler.submit_job("", ctx, &test_plan(2), None).await?;
+
+        // Job submission runs asynchronously through the event loop.
+        let submitted = await_condition(Duration::from_millis(10), 100, || {
+            futures::future::ready(Ok(!reasons.lock().unwrap().is_empty()))
+        })
+        .await?;
+        assert!(submitted, "JobSubmitted callback never fired");
+        assert_eq!(
+            reasons.lock().unwrap().first(),
+            Some(&WorkAvailableReason::JobSubmitted {
+                job_id: job_id.clone()
+            })
+        );
+
+        // Pull the shuffle stage's tasks; the callback promised they are
+        // visible by the time it fired.
+        let exec_meta = ExecutorRegistration {
+            id: "executor-1".to_owned(),
+            host: Some("localhost".to_owned()),
+            port: 50051,
+            grpc_port: 50052,
+            specification: Some(ExecutorSpecification::default().with_vcores(2).into()),
+            os_info: Some(ExecutorOperatingSystemSpecification::default().into()),
+            ballista_protocol_version: BALLISTA_PROTOCOL_VERSION,
+        };
+        let polled = scheduler
+            .poll_work(Request::new(PollWorkParams {
+                metadata: Some(exec_meta.clone()),
+                num_free_vcores: 2,
+                task_status: vec![],
+            }))
+            .await
+            .expect("poll_work failed")
+            .into_inner();
+        assert!(
+            !polled.tasks.is_empty(),
+            "expected tasks after JobSubmitted"
+        );
+
+        // Report the pulled tasks as successful; each writes the plan's two
+        // shuffle output partitions.
+        let task_status = polled
+            .tasks
+            .iter()
+            .map(|task| TaskStatus {
+                task_id: task.task_id,
+                job_id: task.job_id.clone(),
+                stage_id: task.stage_id,
+                stage_attempt_num: task.stage_attempt_num,
+                launch_time: 0,
+                start_exec_time: 0,
+                end_exec_time: 0,
+                metrics: vec![],
+                status: Some(task_status::Status::Successful(SuccessfulTask {
+                    executor_id: exec_meta.id.clone(),
+                    partitions: (0..2)
+                        .map(|partition_id| ShuffleWritePartition {
+                            partition_id,
+                            num_batches: 1,
+                            num_rows: 1,
+                            num_bytes: 1,
+                            file_id: None,
+                            is_sort_shuffle: false,
+                        })
+                        .collect(),
+                    runtime_stats: vec![],
+                })),
+            })
+            .collect();
+        scheduler
+            .poll_work(Request::new(PollWorkParams {
+                metadata: Some(exec_meta),
+                num_free_vcores: 2,
+                task_status,
+            }))
+            .await
+            .expect("poll_work with task status failed");
+
+        // Completing the shuffle stage resolves the final stage.
+        let resolved = await_condition(Duration::from_millis(10), 100, || {
+            futures::future::ready(Ok(reasons.lock().unwrap().contains(
+                &WorkAvailableReason::NewStagesRunnable {
+                    job_id: job_id.clone(),
+                },
+            )))
+        })
+        .await?;
+        assert!(
+            resolved,
+            "expected NewStagesRunnable, got {:?}",
+            reasons.lock().unwrap()
+        );
+
+        Ok(())
+    }
+
     #[tokio::test]
     async fn test_running_job_fails_when_all_executors_are_lost() -> Result<()> {
         let plan = test_plan(10);