1892: feat: allow an external semaphore for executor task concurrency

.cursor/rules.md

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

AGENTS.md

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

CLAUDE.md

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

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<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan, C>(
     mut scheduler: SchedulerGrpcClient<C>,
     executor: Arc<Executor>,
     codec: BallistaCodec<T, U>,
+    available_task_slots: Option<Arc<Semaphore>>,
 ) -> Result<(), BallistaError>
 where
     C: tonic::client::GrpcService<tonic::body::Body>,
@@ -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::<TaskStatus>();

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
             )));
         }
     };

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(())
 }
 

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! {