(datafusion-ballista) branch main updated: feat: allow an external semaphore for executor task concurrency (#1892)

[mi...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

milenkovicm pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/datafusion-ballista.git


The following commit(s) were added to refs/heads/main by this push:
     new 5408c7613 feat: allow an external semaphore for executor task concurrency (#1892)
5408c7613 is described below

commit 5408c7613e01dbc33781d267feabdd212a18ed1a
Author: Jack Eadie <ja...@gmail.com>
AuthorDate: Mon Jul 20 21:42:07 2026 +1000

    feat: allow an external semaphore for executor task concurrency (#1892)
    
    * feat: allow an external semaphore for executor task concurrency
    
    Adds an optional available_task_slots: Option<Arc<Semaphore>> parameter to
    the executor poll_loop. When provided, the caller supplies the semaphore
    that bounds concurrent task execution; this enables sharing it across poll
    loops connected to different schedulers and observing executor busy state
    via available_permits(). When None, the loop creates one internally sized
    to the executor's task slots (unchanged behavior).
    
    Ported from spiceai/datafusion-ballista#14 for upstreaming. Note: layers
    on the poll_loop signature; if landed after #12 (poll_now_notify) it needs
    a trivial rebase to add both parameters.
    
    * fix(examples): pass optional semaphore arg in mtls-cluster and reformat standalone
    
    * fix(executor): address PR review feedback on shared-semaphore poll_loop
    
    - Handle AcquireError from semaphore acquire/acquire_owned by mapping to
      BallistaError instead of panicking, since the semaphore may now be
      externally owned and closed by the caller
    - Assert that the provided semaphore has at least one permit on entry to
      fail fast rather than deadlock
    - Rewrite poll_loop doc comment to be public-facing: remove references
      to internal variable names and implementation types, document shared-
      semaphore over-commit semantics and sizing tradeoffs in behavioral
      terms, and add a # Panics section for the zero-permit case
    
    * style: cargo fmt execution_loop.rs
    
    ---------
    
    Co-authored-by: Phillip LeBlanc <87...@users.noreply.github.com>
---
 ballista/executor/src/execution_loop.rs   | 50 ++++++++++++++++++++++++++-----
 ballista/executor/src/executor_process.rs |  1 +
 ballista/executor/src/standalone.rs       |  1 +
 examples/examples/mtls-cluster.rs         |  2 +-
 4 files changed, 45 insertions(+), 9 deletions(-)

diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs
index da4ff19cd..c4fac99d2 100644
--- a/ballista/executor/src/execution_loop.rs
+++ b/ballista/executor/src/execution_loop.rs
@@ -51,16 +51,38 @@ use tonic::codegen::{Body, Bytes, StdError};
 
 /// Main execution loop that polls the scheduler for available tasks.
 ///
-/// This function runs indefinitely, periodically asking the scheduler for
-/// work. When tasks are received, they are executed on a dedicated thread
-/// pool and results are reported back to the scheduler.
+/// Runs indefinitely, periodically asking the scheduler for work. When tasks
+/// are received they are executed concurrently and results are reported back
+/// to the scheduler.
 ///
-/// The loop respects the executor's concurrent task limit via a semaphore,
-/// ensuring no more than the configured number of tasks run simultaneously.
+/// Concurrency is bounded by a semaphore. Pass `free_vcores` to supply your
+/// own semaphore — useful for sharing a single concurrency limit across
+/// multiple poll loops or for observing executor load from outside.
+/// Pass `None` to have the loop create a semaphore sized to the executor's
+/// configured vcore count.
+///
+/// **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
+/// still caps total concurrent execution — tasks that cannot run immediately
+/// wait for capacity — but both schedulers may over-commit relative to what
+/// the semaphore can actually admit at once. This is intentional: the
+/// semaphore acts as an execution throttle, not a reservation system.
+///
+/// **Semaphore sizing**: if the provided semaphore allows more concurrent
+/// tasks than the executor's thread pool has threads, excess admitted tasks
+/// will queue behind running ones. The caller is responsible for sizing the
+/// semaphore appropriately for their thread pool.
+///
+/// # Panics
+///
+/// Panics on startup if `free_vcores` is a semaphore with zero permits,
+/// which would cause the loop to deadlock immediately.
 pub async fn poll_loop<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan, C>(
     mut scheduler: SchedulerGrpcClient<C>,
     executor: Arc<Executor>,
     codec: BallistaCodec<T, U>,
+    free_vcores: Option<Arc<Semaphore>>,
     health: crate::health::ExecutorHealth,
 ) -> Result<(), BallistaError>
 where
@@ -76,7 +98,13 @@ where
         .unwrap()
         .clone()
         .into();
-    let free_vcores = Arc::new(Semaphore::new(executor_specification.vcores as usize));
+    let free_vcores = free_vcores.unwrap_or_else(|| {
+        Arc::new(Semaphore::new(executor_specification.vcores as usize))
+    });
+    assert!(
+        free_vcores.available_permits() > 0,
+        "free_vcores semaphore must have at least one permit; passing a closed or zero-permit semaphore would deadlock the poll loop"
+    );
 
     let (task_status_sender, mut task_status_receiver) =
         std::sync::mpsc::channel::<TaskStatus>();
@@ -87,7 +115,10 @@ where
 
     loop {
         // Wait for a vcore permit before asking for new work.
-        let permit = free_vcores.acquire().await.unwrap();
+        let permit = free_vcores
+            .acquire()
+            .await
+            .map_err(|_| BallistaError::Internal("vcore semaphore closed".to_string()))?;
         // Make the vcore available again for the actual bind below.
         drop(permit);
 
@@ -137,7 +168,10 @@ where
                     let task_status_sender = task_status_sender.clone();
 
                     // Acquire a vcore permit for the task.
-                    let permit = free_vcores.clone().acquire_owned().await.unwrap();
+                    let permit =
+                        free_vcores.clone().acquire_owned().await.map_err(|_| {
+                            BallistaError::Internal("vcore semaphore closed".to_string())
+                        })?;
 
                     let start_exec_time = SystemTime::now()
                         .duration_since(UNIX_EPOCH)
diff --git a/ballista/executor/src/executor_process.rs b/ballista/executor/src/executor_process.rs
index 9bce0d549..d8f30163b 100644
--- a/ballista/executor/src/executor_process.rs
+++ b/ballista/executor/src/executor_process.rs
@@ -564,6 +564,7 @@ pub async fn start_executor_process(
                 scheduler.clone(),
                 executor.clone(),
                 default_codec,
+                None, // free_vcores: use internal semaphore
                 health,
             )));
         }
diff --git a/ballista/executor/src/standalone.rs b/ballista/executor/src/standalone.rs
index d39841e8b..fd9841a0e 100644
--- a/ballista/executor/src/standalone.rs
+++ b/ballista/executor/src/standalone.rs
@@ -147,6 +147,7 @@ pub async fn new_standalone_executor_from_builder(
         scheduler,
         executor,
         codec,
+        None,
         crate::health::ExecutorHealth::new(),
     ));
     Ok(())
diff --git a/examples/examples/mtls-cluster.rs b/examples/examples/mtls-cluster.rs
index 3533019b9..609935d6b 100644
--- a/examples/examples/mtls-cluster.rs
+++ b/examples/examples/mtls-cluster.rs
@@ -433,7 +433,7 @@ async fn run_executor() -> Result<(), Box<dyn std::error::Error>> {
     info!("Starting execution poll loop...");
     let health = ballista_executor::health::ExecutorHealth::new();
     let poll_handle = tokio::spawn(async move {
-        execution_loop::poll_loop(scheduler, executor, codec, health).await
+        execution_loop::poll_loop(scheduler, executor, codec, None, health).await
     });
 
     tokio::select! {


---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@datafusion.apache.org
For additional commands, e-mail: commits-help@datafusion.apache.org