From 52c78ddacb819ab84a40a5a04c4c236fd3146ac3 Mon Sep 17 00:00:00 2001 From: Jeadie Date: Tue, 23 Jun 2026 14:03:37 +1000 Subject: [PATCH 1/4] feat: allow an external semaphore for executor task concurrency Adds an optional available_task_slots: Option> 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. --- ballista/executor/src/execution_loop.rs | 11 +++++++++-- ballista/executor/src/executor_process.rs | 1 + ballista/executor/src/standalone.rs | 4 +++- 3 files changed, 13 insertions(+), 3 deletions(-) diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs index f64fae3014..bc67540555 100644 --- a/ballista/executor/src/execution_loop.rs +++ b/ballista/executor/src/execution_loop.rs @@ -57,10 +57,16 @@ 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. +/// +/// `available_task_slots`, when provided, lets the caller supply the semaphore +/// controlling task concurrency (e.g. to share it across poll loops or to +/// observe executor busy state via `available_permits()`). When `None`, a +/// semaphore sized to the executor's task slots is created internally. pub async fn poll_loop( mut scheduler: SchedulerGrpcClient, executor: Arc, codec: BallistaCodec, + available_task_slots: Option>, ) -> Result<(), BallistaError> where C: tonic::client::GrpcService, @@ -75,8 +81,9 @@ where .unwrap() .clone() .into(); - let available_task_slots = - Arc::new(Semaphore::new(executor_specification.task_slots as usize)); + let available_task_slots = available_task_slots.unwrap_or_else(|| { + Arc::new(Semaphore::new(executor_specification.task_slots as usize)) + }); let (task_status_sender, mut task_status_receiver) = std::sync::mpsc::channel::(); diff --git a/ballista/executor/src/executor_process.rs b/ballista/executor/src/executor_process.rs index 6050649ab9..36b2dcbb42 100644 --- a/ballista/executor/src/executor_process.rs +++ b/ballista/executor/src/executor_process.rs @@ -505,6 +505,7 @@ pub async fn start_executor_process( scheduler.clone(), executor.clone(), default_codec, + None, // available_task_slots: use internal semaphore ))); } }; diff --git a/ballista/executor/src/standalone.rs b/ballista/executor/src/standalone.rs index bd72dee4b4..f21e5c765f 100644 --- a/ballista/executor/src/standalone.rs +++ b/ballista/executor/src/standalone.rs @@ -142,7 +142,9 @@ 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(()) } From 65bcfc3d976518fc2937fc3c2837d3cdfe34d33e Mon Sep 17 00:00:00 2001 From: jeadie Date: Tue, 23 Jun 2026 15:22:56 +1000 Subject: [PATCH 2/4] fix(examples): pass optional semaphore arg in mtls-cluster and reformat standalone --- ballista/executor/src/standalone.rs | 4 +--- examples/examples/mtls-cluster.rs | 2 +- 2 files changed, 2 insertions(+), 4 deletions(-) diff --git a/ballista/executor/src/standalone.rs b/ballista/executor/src/standalone.rs index f21e5c765f..a77ba1b868 100644 --- a/ballista/executor/src/standalone.rs +++ b/ballista/executor/src/standalone.rs @@ -142,9 +142,7 @@ pub async fn new_standalone_executor_from_builder( )), ); - tokio::spawn(execution_loop::poll_loop( - scheduler, executor, codec, None, - )); + tokio::spawn(execution_loop::poll_loop(scheduler, executor, codec, None)); Ok(()) } diff --git a/examples/examples/mtls-cluster.rs b/examples/examples/mtls-cluster.rs index 593248183e..128e2b64a7 100644 --- a/examples/examples/mtls-cluster.rs +++ b/examples/examples/mtls-cluster.rs @@ -430,7 +430,7 @@ async fn run_executor() -> Result<(), Box> { // 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! { From b75201135e732eedbbc3bba80bc1405fbef734c7 Mon Sep 17 00:00:00 2001 From: jeadie Date: Fri, 3 Jul 2026 09:45:34 +1000 Subject: [PATCH 3/4] 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 --- ballista/executor/src/execution_loop.rs | 54 +++++++++++++++++++------ 1 file changed, 42 insertions(+), 12 deletions(-) diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs index bc67540555..5ece368485 100644 --- a/ballista/executor/src/execution_loop.rs +++ b/ballista/executor/src/execution_loop.rs @@ -51,17 +51,33 @@ 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 `available_task_slots` 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 task-slot count. /// -/// `available_task_slots`, when provided, lets the caller supply the semaphore -/// controlling task concurrency (e.g. to share it across poll loops or to -/// observe executor busy state via `available_permits()`). When `None`, a -/// semaphore sized to the executor's task slots is created internally. +/// **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 `available_task_slots` is a semaphore with zero +/// permits, which would cause the loop to deadlock immediately. pub async fn poll_loop( mut scheduler: SchedulerGrpcClient, executor: Arc, @@ -84,6 +100,10 @@ where let available_task_slots = available_task_slots.unwrap_or_else(|| { Arc::new(Semaphore::new(executor_specification.task_slots as usize)) }); + assert!( + available_task_slots.available_permits() > 0, + "available_task_slots 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::(); @@ -94,7 +114,10 @@ where loop { // Wait for task slots to be available before asking for new work - let permit = available_task_slots.acquire().await.unwrap(); + let permit = available_task_slots + .acquire() + .await + .map_err(|_| BallistaError::Internal("task slot semaphore closed".to_string()))?; // Make the slot available again drop(permit); @@ -140,8 +163,15 @@ where let task_status_sender = task_status_sender.clone(); // Acquire a permit/slot for the task - let permit = - available_task_slots.clone().acquire_owned().await.unwrap(); + let permit = available_task_slots + .clone() + .acquire_owned() + .await + .map_err(|_| { + BallistaError::Internal( + "task slot semaphore closed".to_string(), + ) + })?; let start_exec_time = SystemTime::now() .duration_since(UNIX_EPOCH) From ab34429d1f1fe1c80d1eb2b5896a38d7ac18f0dd Mon Sep 17 00:00:00 2001 From: jeadie Date: Wed, 8 Jul 2026 16:56:20 +1000 Subject: [PATCH 4/4] style: cargo fmt execution_loop.rs --- ballista/executor/src/execution_loop.rs | 24 +++++++++++------------- 1 file changed, 11 insertions(+), 13 deletions(-) diff --git a/ballista/executor/src/execution_loop.rs b/ballista/executor/src/execution_loop.rs index 5ece368485..6ace64e69a 100644 --- a/ballista/executor/src/execution_loop.rs +++ b/ballista/executor/src/execution_loop.rs @@ -114,10 +114,9 @@ where loop { // Wait for task slots to be available before asking for new work - let permit = available_task_slots - .acquire() - .await - .map_err(|_| BallistaError::Internal("task slot semaphore closed".to_string()))?; + let permit = available_task_slots.acquire().await.map_err(|_| { + BallistaError::Internal("task slot semaphore closed".to_string()) + })?; // Make the slot available again drop(permit); @@ -163,15 +162,14 @@ where let task_status_sender = task_status_sender.clone(); // Acquire a permit/slot for the task - let permit = available_task_slots - .clone() - .acquire_owned() - .await - .map_err(|_| { - BallistaError::Internal( - "task slot semaphore closed".to_string(), - ) - })?; + let permit = + available_task_slots.clone().acquire_owned().await.map_err( + |_| { + BallistaError::Internal( + "task slot semaphore closed".to_string(), + ) + }, + )?; let start_exec_time = SystemTime::now() .duration_since(UNIX_EPOCH)