Cluster state refactor part 1

ballista/core/src/serde/generated/ballista.rs

@@ -1985,7 +1985,7 @@ pub mod executor_grpc_client {
 pub mod scheduler_grpc_server {
     #![allow(unused_variables, dead_code, missing_docs, clippy::let_unit_value)]
     use tonic::codegen::*;
-    /// Generated trait containing gRPC methods that should be implemented for use with SchedulerGrpcServer.
+    ///Generated trait containing gRPC methods that should be implemented for use with SchedulerGrpcServer.
     #[async_trait]
     pub trait SchedulerGrpc: Send + Sync + 'static {
         /// Executors must poll the scheduler for heartbeat and to receive tasks
@@ -2531,7 +2531,7 @@ pub mod scheduler_grpc_server {
 pub mod executor_grpc_server {
     #![allow(unused_variables, dead_code, missing_docs, clippy::let_unit_value)]
     use tonic::codegen::*;
-    /// Generated trait containing gRPC methods that should be implemented for use with ExecutorGrpcServer.
+    ///Generated trait containing gRPC methods that should be implemented for use with ExecutorGrpcServer.
     #[async_trait]
     pub trait ExecutorGrpc: Send + Sync + 'static {
         async fn launch_task(

ballista/scheduler/scheduler_config_spec.toml

@@ -36,6 +36,13 @@ type = "ballista_scheduler::state::backend::StateBackend"
 doc = "The configuration backend for the scheduler, possible values: etcd, memory, sled. Default: sled"
 default = "ballista_scheduler::state::backend::StateBackend::Sled"
 
+[[param]]
+abbr = "c"
+name = "cluster_backend"
+type = "ballista_scheduler::state::backend::StateBackend"
+doc = "The configuration backend for the scheduler cluster state, possible values: etcd, memory, sled. Default: sled"
+default = "ballista_scheduler::state::backend::StateBackend::Sled"
+
 [[param]]
 abbr = "n"
 name = "namespace"

ballista/scheduler/src/bin/main.rs

@@ -45,6 +45,7 @@ mod config {
 
 use ballista_core::config::LogRotationPolicy;
 use ballista_scheduler::config::SchedulerConfig;
+use ballista_scheduler::state::backend::cluster::DefaultClusterState;
 use config::prelude::*;
 use tracing_subscriber::EnvFilter;
 
@@ -60,6 +61,16 @@ async fn main() -> Result<()> {
         std::process::exit(0);
     }
 
+    let config_backend = init_kv_backend(&opt.config_backend, &opt).await?;
+
+    let cluster_state = if opt.cluster_backend == opt.config_backend {
+        Arc::new(DefaultClusterState::new(config_backend.clone()))
+    } else {
+        let cluster_kv_store = init_kv_backend(&opt.cluster_backend, &opt).await?;
+
+        Arc::new(DefaultClusterState::new(cluster_kv_store))
+    };
+
     let special_mod_log_level = opt.log_level_setting;
     let namespace = opt.namespace;
     let external_host = opt.external_host;
@@ -110,13 +121,31 @@ async fn main() -> Result<()> {
     let addr = format!("{}:{}", bind_host, port);
     let addr = addr.parse()?;
 
-    let config_backend: Arc<dyn StateBackendClient> = match opt.config_backend {
+    let config = SchedulerConfig {
+        scheduling_policy: opt.scheduler_policy,
+        event_loop_buffer_size: opt.event_loop_buffer_size,
+        executor_slots_policy: opt.executor_slots_policy,
+        finished_job_data_clean_up_interval_seconds: opt
+            .finished_job_data_clean_up_interval_seconds,
+        finished_job_state_clean_up_interval_seconds: opt
+            .finished_job_state_clean_up_interval_seconds,
+        advertise_flight_sql_endpoint: opt.advertise_flight_sql_endpoint,
+    };
+    start_server(scheduler_name, config_backend, cluster_state, addr, config).await?;
+    Ok(())
+}
+
+async fn init_kv_backend(
+    backend: &StateBackend,
+    opt: &Config,
+) -> Result<Arc<dyn StateBackendClient>> {
+    let cluster_backend: Arc<dyn StateBackendClient> = match backend {
         #[cfg(feature = "etcd")]
         StateBackend::Etcd => {
-            let etcd = etcd_client::Client::connect(&[opt.etcd_urls], None)
+            let etcd = etcd_client::Client::connect(&[opt.etcd_urls.clone()], None)
                 .await
                 .context("Could not connect to etcd")?;
-            Arc::new(EtcdClient::new(namespace.clone(), etcd))
+            Arc::new(EtcdClient::new(opt.namespace.clone(), etcd))
         }
         #[cfg(not(feature = "etcd"))]
         StateBackend::Etcd => {
@@ -134,7 +163,7 @@ async fn main() -> Result<()> {
             } else {
                 println!("{}", opt.sled_dir);
                 Arc::new(
-                    SledClient::try_new(opt.sled_dir)
+                    SledClient::try_new(opt.sled_dir.clone())
                         .context("Could not create sled config backend")?,
                 )
             }
@@ -148,16 +177,5 @@ async fn main() -> Result<()> {
         StateBackend::Memory => Arc::new(MemoryBackendClient::new()),
     };
 
-    let config = SchedulerConfig {
-        scheduling_policy: opt.scheduler_policy,
-        event_loop_buffer_size: opt.event_loop_buffer_size,
-        executor_slots_policy: opt.executor_slots_policy,
-        finished_job_data_clean_up_interval_seconds: opt
-            .finished_job_data_clean_up_interval_seconds,
-        finished_job_state_clean_up_interval_seconds: opt
-            .finished_job_state_clean_up_interval_seconds,
-        advertise_flight_sql_endpoint: opt.advertise_flight_sql_endpoint,
-    };
-    start_server(scheduler_name, config_backend, addr, config).await?;
-    Ok(())
+    Ok(cluster_backend)
 }

ballista/scheduler/src/config.rs

@@ -90,6 +90,11 @@ impl SchedulerConfig {
         self.advertise_flight_sql_endpoint = endpoint;
         self
     }
+
+    pub fn with_executor_slots_policy(mut self, policy: SlotsPolicy) -> Self {
+        self.executor_slots_policy = policy;
+        self
+    }
 }
 
 // an enum used to configure the executor slots policy

ballista/scheduler/src/scheduler_process.rs

@@ -41,11 +41,13 @@ use crate::flight_sql::FlightSqlServiceImpl;
 use crate::metrics::default_metrics_collector;
 use crate::scheduler_server::externalscaler::external_scaler_server::ExternalScalerServer;
 use crate::scheduler_server::SchedulerServer;
+use crate::state::backend::cluster::ClusterState;
 use crate::state::backend::StateBackendClient;
 
 pub async fn start_server(
     scheduler_name: String,
     config_backend: Arc<dyn StateBackendClient>,
+    cluster_state: Arc<dyn ClusterState>,
     addr: SocketAddr,
     config: SchedulerConfig,
 ) -> Result<()> {
@@ -65,6 +67,7 @@ pub async fn start_server(
         SchedulerServer::new(
             scheduler_name,
             config_backend.clone(),
+            cluster_state,
             BallistaCodec::default(),
             config,
             metrics_collector,

ballista/scheduler/src/scheduler_server/grpc.rs

@@ -588,6 +588,7 @@ mod test {
     use ballista_core::serde::BallistaCodec;
     use ballista_core::utils::default_session_builder;
 
+    use crate::state::backend::cluster::DefaultClusterState;
     use crate::state::executor_manager::DEFAULT_EXECUTOR_TIMEOUT_SECONDS;
     use crate::state::{backend::sled::SledClient, SchedulerState};
 
@@ -596,10 +597,12 @@ mod test {
     #[tokio::test]
     async fn test_poll_work() -> Result<(), BallistaError> {
         let state_storage = Arc::new(SledClient::try_new_temporary()?);
+        let cluster_state = Arc::new(DefaultClusterState::new(state_storage.clone()));
         let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
             SchedulerServer::new(
                 "localhost:50050".to_owned(),
                 state_storage.clone(),
+                cluster_state.clone(),
                 BallistaCodec::default(),
                 SchedulerConfig::default(),
                 default_metrics_collector().unwrap(),
@@ -627,6 +630,7 @@ mod test {
         let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
             SchedulerState::new_with_default_scheduler_name(
                 state_storage.clone(),
+                cluster_state.clone(),
                 default_session_builder,
                 BallistaCodec::default(),
             );
@@ -660,6 +664,7 @@ mod test {
         let state: SchedulerState<LogicalPlanNode, PhysicalPlanNode> =
             SchedulerState::new_with_default_scheduler_name(
                 state_storage.clone(),
+                cluster_state,
                 default_session_builder,
                 BallistaCodec::default(),
             );
@@ -683,10 +688,12 @@ mod test {
     #[tokio::test]
     async fn test_stop_executor() -> Result<(), BallistaError> {
         let state_storage = Arc::new(SledClient::try_new_temporary()?);
+        let cluster_state = Arc::new(DefaultClusterState::new(state_storage.clone()));
         let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
             SchedulerServer::new(
                 "localhost:50050".to_owned(),
-                state_storage.clone(),
+                state_storage,
+                cluster_state,
                 BallistaCodec::default(),
                 SchedulerConfig::default(),
                 default_metrics_collector().unwrap(),
@@ -764,10 +771,12 @@ mod test {
     #[ignore]
     async fn test_expired_executor() -> Result<(), BallistaError> {
         let state_storage = Arc::new(SledClient::try_new_temporary()?);
+        let cluster_state = Arc::new(DefaultClusterState::new(state_storage.clone()));
         let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
             SchedulerServer::new(
                 "localhost:50050".to_owned(),
-                state_storage.clone(),
+                state_storage,
+                cluster_state,
                 BallistaCodec::default(),
                 SchedulerConfig::default(),
                 default_metrics_collector().unwrap(),

ballista/scheduler/src/scheduler_server/mod.rs

@@ -35,6 +35,7 @@ use log::{error, warn};
 
 use crate::scheduler_server::event::QueryStageSchedulerEvent;
 use crate::scheduler_server::query_stage_scheduler::QueryStageScheduler;
+use crate::state::backend::cluster::ClusterState;
 use crate::state::backend::StateBackendClient;
 use crate::state::executor_manager::{
     ExecutorManager, ExecutorReservation, DEFAULT_EXECUTOR_TIMEOUT_SECONDS,
@@ -69,12 +70,14 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
     pub fn new(
         scheduler_name: String,
         config_backend: Arc<dyn StateBackendClient>,
+        cluster_state: Arc<dyn ClusterState>,
         codec: BallistaCodec<T, U>,
         config: SchedulerConfig,
         metrics_collector: Arc<dyn SchedulerMetricsCollector>,
     ) -> Self {
         let state = Arc::new(SchedulerState::new(
             config_backend,
+            cluster_state,
             default_session_builder,
             codec,
             scheduler_name.clone(),
@@ -97,17 +100,53 @@ impl<T: 'static + AsLogicalPlan, U: 'static + AsExecutionPlan> SchedulerServer<T
         }
     }
 
+    pub fn with_session_builder(
+        scheduler_name: String,
+        config_backend: Arc<dyn StateBackendClient>,
+        cluster_backend: Arc<dyn ClusterState>,
+        codec: BallistaCodec<T, U>,
+        config: SchedulerConfig,
+        session_builder: SessionBuilder,
+        metrics_collector: Arc<dyn SchedulerMetricsCollector>,
+    ) -> Self {
+        let state = Arc::new(SchedulerState::new(
+            config_backend,
+            cluster_backend,
+            session_builder,
+            codec,
+            scheduler_name.clone(),
+            config.clone(),
+        ));
+        let query_stage_scheduler =
+            Arc::new(QueryStageScheduler::new(state.clone(), metrics_collector));
+        let query_stage_event_loop = EventLoop::new(
+            "query_stage".to_owned(),
+            config.event_loop_buffer_size as usize,
+            query_stage_scheduler.clone(),
+        );
+
+        Self {
+            scheduler_name,
+            start_time: timestamp_millis() as u128,
+            state,
+            query_stage_event_loop,
+            query_stage_scheduler,
+        }
+    }
+
     #[allow(dead_code)]
     pub(crate) fn with_task_launcher(
         scheduler_name: String,
         config_backend: Arc<dyn StateBackendClient>,
+        cluster_backend: Arc<dyn ClusterState>,
         codec: BallistaCodec<T, U>,
         config: SchedulerConfig,
         metrics_collector: Arc<dyn SchedulerMetricsCollector>,
         task_launcher: Arc<dyn TaskLauncher>,
     ) -> Self {
         let state = Arc::new(SchedulerState::with_task_launcher(
             config_backend,
+            cluster_backend,
             default_session_builder,
             codec,
             scheduler_name.clone(),
@@ -330,6 +369,7 @@ mod test {
     use ballista_core::serde::BallistaCodec;
 
     use crate::scheduler_server::{timestamp_millis, SchedulerServer};
+    use crate::state::backend::cluster::DefaultClusterState;
     use crate::state::backend::sled::SledClient;
 
     use crate::test_utils::{
@@ -599,10 +639,12 @@ mod test {
         scheduling_policy: TaskSchedulingPolicy,
     ) -> Result<SchedulerServer<LogicalPlanNode, PhysicalPlanNode>> {
         let state_storage = Arc::new(SledClient::try_new_temporary()?);
+        let cluster_state = Arc::new(DefaultClusterState::new(state_storage.clone()));
         let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
             SchedulerServer::new(
                 "localhost:50050".to_owned(),
-                state_storage.clone(),
+                state_storage,
+                cluster_state,
                 BallistaCodec::default(),
                 SchedulerConfig::default().with_scheduler_policy(scheduling_policy),
                 Arc::new(TestMetricsCollector::default()),

ballista/scheduler/src/standalone.rs

@@ -17,6 +17,7 @@
 
 use crate::config::SchedulerConfig;
 use crate::metrics::default_metrics_collector;
+use crate::state::backend::cluster::DefaultClusterState;
 use crate::{scheduler_server::SchedulerServer, state::backend::sled::SledClient};
 use ballista_core::serde::protobuf::PhysicalPlanNode;
 use ballista_core::serde::BallistaCodec;
@@ -31,14 +32,15 @@ use std::{net::SocketAddr, sync::Arc};
 use tokio::net::TcpListener;
 
 pub async fn new_standalone_scheduler() -> Result<SocketAddr> {
-    let client = SledClient::try_new_temporary()?;
+    let backend = Arc::new(SledClient::try_new_temporary()?);
 
     let metrics_collector = default_metrics_collector()?;
 
     let mut scheduler_server: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
         SchedulerServer::new(
             "localhost:50050".to_owned(),
-            Arc::new(client),
+            backend.clone(),
+            Arc::new(DefaultClusterState::new(backend)),
             BallistaCodec::default(),
             SchedulerConfig::default(),
             metrics_collector,