Provide a memory StateBackendClient (#523)

* Rename StateBackend::Standalone to StateBackend:Sled

* Copy utility files from sled crate since they cannot be used directly

* Provide a memory StateBackendClient

* Fix dashmap deadlock issue

* Fix for the comments

Co-authored-by: yangzhong <[email protected]>

Author: yahoNanJing

ballista/scheduler/scheduler_config_spec.toml

@@ -33,8 +33,8 @@ doc = "Route for proxying flight results via scheduler. Should be of the form 'I
 abbr = "b"
 name = "config_backend"
 type = "ballista_scheduler::state::backend::StateBackend"
-doc = "The configuration backend for the scheduler, possible values: etcd, standalone. Default: standalone"
-default = "ballista_scheduler::state::backend::StateBackend::Standalone"
+doc = "The configuration backend for the scheduler, possible values: etcd, memory, sled. Default: sled"
+default = "ballista_scheduler::state::backend::StateBackend::Sled"
 
 [[param]]
 abbr = "n"

ballista/scheduler/src/main.rs

@@ -21,6 +21,7 @@ use anyhow::{Context, Result};
 #[cfg(feature = "flight-sql")]
 use arrow_flight::flight_service_server::FlightServiceServer;
 use ballista_scheduler::scheduler_server::externalscaler::external_scaler_server::ExternalScalerServer;
+use ballista_scheduler::state::backend::memory::MemoryBackendClient;
 use futures::future::{self, Either, TryFutureExt};
 use hyper::{server::conn::AddrStream, service::make_service_fn, Server};
 use std::convert::Infallible;
@@ -37,7 +38,7 @@ use ballista_scheduler::api::{get_routes, EitherBody, Error};
 #[cfg(feature = "etcd")]
 use ballista_scheduler::state::backend::etcd::EtcdClient;
 #[cfg(feature = "sled")]
-use ballista_scheduler::state::backend::standalone::StandaloneClient;
+use ballista_scheduler::state::backend::sled::SledClient;
 use datafusion_proto::protobuf::LogicalPlanNode;
 
 use ballista_scheduler::scheduler_server::SchedulerServer;
@@ -211,10 +212,6 @@ async fn main() -> Result<()> {
     let addr = addr.parse()?;
 
     let config_backend: Arc<dyn StateBackendClient> = match opt.config_backend {
-        #[cfg(not(any(feature = "sled", feature = "etcd")))]
-        _ => std::compile_error!(
-            "To build the scheduler enable at least one config backend feature (`etcd` or `sled`)"
-        ),
         #[cfg(feature = "etcd")]
         StateBackend::Etcd => {
             let etcd = etcd_client::Client::connect(&[opt.etcd_urls], None)
@@ -229,26 +226,27 @@ async fn main() -> Result<()> {
             )
         }
         #[cfg(feature = "sled")]
-        StateBackend::Standalone => {
+        StateBackend::Sled => {
             if opt.sled_dir.is_empty() {
                 Arc::new(
-                    StandaloneClient::try_new_temporary()
-                        .context("Could not create standalone config backend")?,
+                    SledClient::try_new_temporary()
+                        .context("Could not create sled config backend")?,
                 )
             } else {
                 println!("{}", opt.sled_dir);
                 Arc::new(
-                    StandaloneClient::try_new(opt.sled_dir)
-                        .context("Could not create standalone config backend")?,
+                    SledClient::try_new(opt.sled_dir)
+                        .context("Could not create sled config backend")?,
                 )
             }
         }
         #[cfg(not(feature = "sled"))]
-        StateBackend::Standalone => {
+        StateBackend::Sled => {
             unimplemented!(
-                "build the scheduler with the `sled` feature to use the standalone config backend"
+                "build the scheduler with the `sled` feature to use the sled config backend"
             )
         }
+        StateBackend::Memory => Arc::new(MemoryBackendClient::new()),
     };
 
     let config = SchedulerConfig {

ballista/scheduler/src/scheduler_server/grpc.rs

@@ -587,13 +587,13 @@ mod test {
     use ballista_core::utils::default_session_builder;
 
     use crate::state::executor_manager::DEFAULT_EXECUTOR_TIMEOUT_SECONDS;
-    use crate::state::{backend::standalone::StandaloneClient, SchedulerState};
+    use crate::state::{backend::sled::SledClient, SchedulerState};
 
     use super::{SchedulerGrpc, SchedulerServer};
 
     #[tokio::test]
     async fn test_poll_work() -> Result<(), BallistaError> {
-        let state_storage = Arc::new(StandaloneClient::try_new_temporary()?);
+        let state_storage = Arc::new(SledClient::try_new_temporary()?);
         let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
             SchedulerServer::new(
                 "localhost:50050".to_owned(),
@@ -680,7 +680,7 @@ mod test {
 
     #[tokio::test]
     async fn test_stop_executor() -> Result<(), BallistaError> {
-        let state_storage = Arc::new(StandaloneClient::try_new_temporary()?);
+        let state_storage = Arc::new(SledClient::try_new_temporary()?);
         let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
             SchedulerServer::new(
                 "localhost:50050".to_owned(),
@@ -761,7 +761,7 @@ mod test {
     #[tokio::test]
     #[ignore]
     async fn test_expired_executor() -> Result<(), BallistaError> {
-        let state_storage = Arc::new(StandaloneClient::try_new_temporary()?);
+        let state_storage = Arc::new(SledClient::try_new_temporary()?);
         let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
             SchedulerServer::new(
                 "localhost:50050".to_owned(),

ballista/scheduler/src/scheduler_server/mod.rs

@@ -330,7 +330,7 @@ mod test {
     use ballista_core::serde::BallistaCodec;
 
     use crate::scheduler_server::{timestamp_millis, SchedulerServer};
-    use crate::state::backend::standalone::StandaloneClient;
+    use crate::state::backend::sled::SledClient;
 
     use crate::test_utils::{
         assert_completed_event, assert_failed_event, assert_no_submitted_event,
@@ -598,7 +598,7 @@ mod test {
     async fn test_scheduler(
         scheduling_policy: TaskSchedulingPolicy,
     ) -> Result<SchedulerServer<LogicalPlanNode, PhysicalPlanNode>> {
-        let state_storage = Arc::new(StandaloneClient::try_new_temporary()?);
+        let state_storage = Arc::new(SledClient::try_new_temporary()?);
         let mut scheduler: SchedulerServer<LogicalPlanNode, PhysicalPlanNode> =
             SchedulerServer::new(
                 "localhost:50050".to_owned(),

ballista/scheduler/src/standalone.rs

@@ -17,9 +17,7 @@
 
 use crate::config::SchedulerConfig;
 use crate::metrics::default_metrics_collector;
-use crate::{
-    scheduler_server::SchedulerServer, state::backend::standalone::StandaloneClient,
-};
+use crate::{scheduler_server::SchedulerServer, state::backend::sled::SledClient};
 use ballista_core::serde::protobuf::PhysicalPlanNode;
 use ballista_core::serde::BallistaCodec;
 use ballista_core::utils::create_grpc_server;
@@ -33,7 +31,7 @@ use std::{net::SocketAddr, sync::Arc};
 use tokio::net::TcpListener;
 
 pub async fn new_standalone_scheduler() -> Result<SocketAddr> {
-    let client = StandaloneClient::try_new_temporary()?;
+    let client = SledClient::try_new_temporary()?;
 
     let metrics_collector = default_metrics_collector()?;
 

ballista/scheduler/src/state/backend/etcd.rs

@@ -34,7 +34,7 @@ use crate::state::backend::{
     Keyspace, Lock, Operation, StateBackendClient, Watch, WatchEvent,
 };
 
-/// A [`StateBackendClient`] implementation that uses etcd to save cluster configuration.
+/// A [`StateBackendClient`] implementation that uses etcd to save cluster state.
 #[derive(Clone)]
 pub struct EtcdClient {
     namespace: String,