executor poll_loop to support oneshot channel for readiness reporting (#2)

Author: mach-kernel

ballista/executor/src/execution_loop.rs

@@ -36,19 +36,22 @@ use datafusion_proto::physical_plan::AsExecutionPlan;
 use futures::FutureExt;
 use log::{debug, error, info, warn};
 use std::any::Any;
+use std::cell::LazyCell;
 use std::convert::TryInto;
 use std::error::Error;
 use std::ops::Deref;
 use std::sync::mpsc::{Receiver, Sender, TryRecvError};
 use std::time::{SystemTime, UNIX_EPOCH};
 use std::{sync::Arc, time::Duration};
+use tokio::sync::oneshot::Sender as OneShotSender;
 use tokio::sync::{OwnedSemaphorePermit, Semaphore};
 use tonic::transport::Channel;
 
 pub async fn poll_loop<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>(
     mut scheduler: SchedulerGrpcClient<Channel>,
     executor: Arc<Executor>,
     codec: BallistaCodec<T, U>,
+    readiness: Option<OneShotSender<String>>,
 ) -> Result<(), BallistaError> {
     let executor_specification: ExecutorSpecification = executor
         .metadata
@@ -67,6 +70,13 @@ pub async fn poll_loop<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
     let dedicated_executor =
         DedicatedExecutor::new("task_runner", executor_specification.task_slots as usize);
 
+    let report_ready = LazyCell::new(|| {
+        if let Some(chan) = readiness {
+            chan.send(executor.metadata.id.clone())
+                .expect("Must send readiness")
+        }
+    });
+
     loop {
         // Wait for task slots to be available before asking for new work
         let permit = available_task_slots.acquire().await.unwrap();
@@ -89,6 +99,8 @@ pub async fn poll_loop<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan>
                 })
                 .await;
 
+        *report_ready;
+
         match poll_work_result {
             Ok(result) => {
                 let PollWorkResult {

ballista/executor/src/executor_process.rs

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

ballista/executor/src/standalone.rs

@@ -131,7 +131,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(())
 }