Make the scheduler event loop buffer size configurable

Author: kyotoYaho

ballista/scheduler/scheduler_config_spec.toml

@@ -72,6 +72,12 @@ type = "ballista_core::config::TaskSchedulingPolicy"
 doc = "The scheduing policy for the scheduler, possible values: pull-staged, push-staged. Default: pull-staged"
 default = "ballista_core::config::TaskSchedulingPolicy::PullStaged"
 
+[[param]]
+name = "event_loop_buffer_size"
+type = "u32"
+default = "10000"
+doc = "Event loop buffer size. Default: 10000"
+
 [[param]]
 name = "executor_slots_policy"
 type = "ballista_scheduler::config::SlotsPolicy"

ballista/scheduler/src/main.rs

@@ -74,6 +74,7 @@ async fn start_server(
     addr: SocketAddr,
     scheduling_policy: TaskSchedulingPolicy,
     slots_policy: SlotsPolicy,
+    event_loop_buffer_size: usize,
 ) -> Result<()> {
     info!(
         "Ballista v{} Scheduler listening on {:?}",
@@ -93,11 +94,13 @@ async fn start_server(
                 slots_policy,
                 BallistaCodec::default(),
                 default_session_builder,
+                event_loop_buffer_size,
             ),
             _ => SchedulerServer::new(
                 scheduler_name,
                 config_backend.clone(),
                 BallistaCodec::default(),
+                event_loop_buffer_size,
             ),
         };
 
@@ -244,12 +247,14 @@ async fn main() -> Result<()> {
 
     let scheduling_policy: TaskSchedulingPolicy = opt.scheduler_policy;
     let slots_policy: SlotsPolicy = opt.executor_slots_policy;
+    let event_loop_buffer_size = opt.event_loop_buffer_size as usize;
     start_server(
         scheduler_name,
         client,
         addr,
         scheduling_policy,
         slots_policy,
+        event_loop_buffer_size,
     )
     .await?;
     Ok(())

ballista/scheduler/src/scheduler_server/grpc.rs

@@ -582,6 +582,7 @@ mod test {
                 "localhost:50050".to_owned(),
                 state_storage.clone(),
                 BallistaCodec::default(),
+                10000,
             );
         scheduler.init().await?;
         let exec_meta = ExecutorRegistration {
@@ -667,6 +668,7 @@ mod test {
                 "localhost:50050".to_owned(),
                 state_storage.clone(),
                 BallistaCodec::default(),
+                10000,
             );
         scheduler.init().await?;
 
@@ -746,6 +748,7 @@ mod test {
                 "localhost:50050".to_owned(),
                 state_storage.clone(),
                 BallistaCodec::default(),
+                10000,
             );
         scheduler.init().await?;
 

ballista/scheduler/src/scheduler_server/mod.rs

@@ -68,6 +68,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
         scheduler_name: String,
         config: Arc<dyn StateBackendClient>,
         codec: BallistaCodec<T, U>,
+        event_loop_buffer_size: usize,
     ) -> Self {
         SchedulerServer::new_with_policy(
             scheduler_name,
@@ -76,6 +77,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             SlotsPolicy::Bias,
             codec,
             default_session_builder,
+            event_loop_buffer_size,
         )
     }
 
@@ -84,6 +86,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
         config: Arc<dyn StateBackendClient>,
         codec: BallistaCodec<T, U>,
         session_builder: SessionBuilder,
+        event_loop_buffer_size: usize,
     ) -> Self {
         SchedulerServer::new_with_policy(
             scheduler_name,
@@ -92,6 +95,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             SlotsPolicy::Bias,
             codec,
             session_builder,
+            event_loop_buffer_size,
         )
     }
 
@@ -102,6 +106,7 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
         slots_policy: SlotsPolicy,
         codec: BallistaCodec<T, U>,
         session_builder: SessionBuilder,
+        event_loop_buffer_size: usize,
     ) -> Self {
         let state = Arc::new(SchedulerState::new(
             config,
@@ -111,18 +116,27 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
             slots_policy,
         ));
 
-        SchedulerServer::new_with_state(scheduler_name, scheduling_policy, state)
+        SchedulerServer::new_with_state(
+            scheduler_name,
+            scheduling_policy,
+            state,
+            event_loop_buffer_size,
+        )
     }
 
     pub(crate) fn new_with_state(
         scheduler_name: String,
         policy: TaskSchedulingPolicy,
         state: Arc<SchedulerState<T, U>>,
+        event_loop_buffer_size: usize,
     ) -> Self {
         let query_stage_scheduler =
             Arc::new(QueryStageScheduler::new(state.clone(), policy));
-        let query_stage_event_loop =
-            EventLoop::new("query_stage".to_owned(), 10000, query_stage_scheduler);
+        let query_stage_event_loop = EventLoop::new(
+            "query_stage".to_owned(),
+            event_loop_buffer_size,
+            query_stage_scheduler,
+        );
         Self {
             scheduler_name,
             state,
@@ -770,6 +784,7 @@ mod test {
                 SlotsPolicy::Bias,
                 BallistaCodec::default(),
                 default_session_builder,
+                10000,
             );
         scheduler.init().await?;
 
@@ -789,6 +804,7 @@ mod test {
                 "localhost:50050".to_owned(),
                 TaskSchedulingPolicy::PushStaged,
                 state,
+                10000,
             );
         scheduler.init().await?;
 

ballista/scheduler/src/standalone.rs

@@ -38,6 +38,7 @@ pub async fn new_standalone_scheduler() -> Result<SocketAddr> {
             "localhost:50050".to_owned(),
             Arc::new(client),
             BallistaCodec::default(),
+            100000,
         );
     scheduler_server.init().await?;
     let server = SchedulerGrpcServer::new(scheduler_server.clone());