feat: add load_based selector in meta

config/metasrv.example.toml

@@ -2,3 +2,5 @@ bind_addr = '127.0.0.1:3002'
 server_addr = '127.0.0.1:3002'
 store_addr = '127.0.0.1:2379'
 datanode_lease_secs = 15
+# selector: 'LeaseBased', 'LoadBased'
+selector = 'LeaseBased'

src/cmd/src/error.rs

@@ -61,6 +61,13 @@ pub enum Error {
         #[snafu(backtrace)]
         source: servers::auth::Error,
     },
+
+    #[snafu(display("Unsupported selector type, {} source: {}", selector_type, source))]
+    UnsupportedSelectorType {
+        selector_type: String,
+        #[snafu(backtrace)]
+        source: meta_srv::error::Error,
+    },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -71,6 +78,7 @@ impl ErrorExt for Error {
             Error::StartDatanode { source } => source.status_code(),
             Error::StartFrontend { source } => source.status_code(),
             Error::StartMetaServer { source } => source.status_code(),
+            Error::UnsupportedSelectorType { source, .. } => source.status_code(),
             Error::ReadConfig { .. } | Error::ParseConfig { .. } | Error::MissingConfig { .. } => {
                 StatusCode::InvalidArguments
             }

src/cmd/src/metasrv.rs

@@ -13,7 +13,7 @@
 // limitations under the License.
 
 use clap::Parser;
-use common_telemetry::logging;
+use common_telemetry::{info, logging};
 use meta_srv::bootstrap;
 use meta_srv::metasrv::MetaSrvOptions;
 use snafu::ResultExt;
@@ -56,6 +56,8 @@ struct StartCommand {
     store_addr: Option<String>,
     #[clap(short, long)]
     config_file: Option<String>,
+    #[clap(short, long)]
+    selector: Option<String>,
 }
 
 impl StartCommand {
@@ -91,13 +93,21 @@ impl TryFrom<StartCommand> for MetaSrvOptions {
         if let Some(addr) = cmd.store_addr {
             opts.store_addr = addr;
         }
+        if let Some(selector_type) = &cmd.selector {
+            opts.selector = selector_type[..]
+                .try_into()
+                .context(error::UnsupportedSelectorTypeSnafu { selector_type })?;
+            info!("Using {} selector", selector_type);
+        }
 
         Ok(opts)
     }
 }
 
 #[cfg(test)]
 mod tests {
+    use meta_srv::selector::SelectorType;
+
     use super::*;
 
     #[test]
@@ -107,11 +117,13 @@ mod tests {
             server_addr: Some("127.0.0.1:3002".to_string()),
             store_addr: Some("127.0.0.1:2380".to_string()),
             config_file: None,
+            selector: Some("LoadBased".to_string()),
         };
         let options: MetaSrvOptions = cmd.try_into().unwrap();
         assert_eq!("127.0.0.1:3002".to_string(), options.bind_addr);
         assert_eq!("127.0.0.1:3002".to_string(), options.server_addr);
         assert_eq!("127.0.0.1:2380".to_string(), options.store_addr);
+        assert_eq!(SelectorType::LoadBased, options.selector);
     }
 
     #[test]
@@ -120,6 +132,7 @@ mod tests {
             bind_addr: None,
             server_addr: None,
             store_addr: None,
+            selector: None,
             config_file: Some(format!(
                 "{}/../../config/metasrv.example.toml",
                 std::env::current_dir().unwrap().as_path().to_str().unwrap()
@@ -130,5 +143,6 @@ mod tests {
         assert_eq!("127.0.0.1:3002".to_string(), options.server_addr);
         assert_eq!("127.0.0.1:2379".to_string(), options.store_addr);
         assert_eq!(15, options.datanode_lease_secs);
+        assert_eq!(SelectorType::LeaseBased, options.selector);
     }
 }

src/meta-srv/Cargo.toml

@@ -8,6 +8,7 @@ license.workspace = true
 mock = []
 
 [dependencies]
+anymap = "1.0.0-beta.2"
 api = { path = "../api" }
 async-trait = "0.1"
 catalog = { path = "../catalog" }

src/meta-srv/src/bootstrap.rs

@@ -60,7 +60,8 @@ pub fn router(meta_srv: MetaSrv) -> Router {
 pub async fn make_meta_srv(opts: MetaSrvOptions) -> Result<MetaSrv> {
     let kv_store = EtcdStore::with_endpoints([&opts.store_addr]).await?;
     let election = EtcdElection::with_endpoints(&opts.server_addr, [&opts.store_addr]).await?;
-    let meta_srv = MetaSrv::new(opts, kv_store, None, Some(election), None).await;
+    let selector = opts.selector.clone().into();
+    let meta_srv = MetaSrv::new(opts, kv_store, Some(selector), Some(election), None).await;
     meta_srv.start().await;
     Ok(meta_srv)
 }

src/meta-srv/src/error.rs

@@ -55,12 +55,33 @@ pub enum Error {
     #[snafu(display("Invalid datanode lease key: {}", key))]
     InvalidLeaseKey { key: String, backtrace: Backtrace },
 
+    #[snafu(display("Invalid datanode stat key: {}", key))]
+    InvalidStatKey { key: String, backtrace: Backtrace },
+
     #[snafu(display("Failed to parse datanode lease key from utf8: {}", source))]
     LeaseKeyFromUtf8 {
         source: std::string::FromUtf8Error,
         backtrace: Backtrace,
     },
 
+    #[snafu(display("Failed to parse datanode lease value from utf8: {}", source))]
+    LeaseValueFromUtf8 {
+        source: std::string::FromUtf8Error,
+        backtrace: Backtrace,
+    },
+
+    #[snafu(display("Failed to parse datanode stat key from utf8: {}", source))]
+    StatKeyFromUtf8 {
+        source: std::string::FromUtf8Error,
+        backtrace: Backtrace,
+    },
+
+    #[snafu(display("Failed to parse datanode stat value from utf8: {}", source))]
+    StatValueFromUtf8 {
+        source: std::string::FromUtf8Error,
+        backtrace: Backtrace,
+    },
+
     #[snafu(display("Failed to serialize to json: {}", input))]
     SerializeToJson {
         input: String,
@@ -132,6 +153,12 @@ pub enum Error {
         key
     ))]
     MoveValue { key: String, backtrace: Backtrace },
+
+    #[snafu(display("Unsupported selector type, {}", selector_type))]
+    UnsupportedSelectorType {
+        selector_type: String,
+        backtrace: Backtrace,
+    },
 }
 
 pub type Result<T> = std::result::Result<T, Error>;
@@ -165,9 +192,14 @@ impl ErrorExt for Error {
             Error::EmptyKey { .. }
             | Error::EmptyTableName { .. }
             | Error::InvalidLeaseKey { .. }
+            | Error::InvalidStatKey { .. }
             | Error::ParseNum { .. }
+            | Error::UnsupportedSelectorType { .. }
             | Error::InvalidArguments { .. } => StatusCode::InvalidArguments,
             Error::LeaseKeyFromUtf8 { .. }
+            | Error::LeaseValueFromUtf8 { .. }
+            | Error::StatKeyFromUtf8 { .. }
+            | Error::StatValueFromUtf8 { .. }
             | Error::UnexceptedSequenceValue { .. }
             | Error::TableRouteNotFound { .. }
             | Error::NextSequence { .. }

src/meta-srv/src/handler.rs

@@ -23,7 +23,7 @@ mod check_leader_handler;
 mod collect_stats_handler;
 mod instruction;
 mod keep_lease_handler;
-pub(crate) mod node_stat;
+pub mod node_stat;
 mod on_leader_start;
 mod persist_stats_handler;
 mod response_header_handler;

src/meta-srv/src/handler/node_stat.rs

@@ -33,7 +33,7 @@ pub struct Stat {
     /// How many tables on this node
     pub table_num: i64,
     /// How many regions on this node
-    pub region_num: i64,
+    pub region_num: Option<u64>,
     pub cpu_usage: f64,
     pub load: f64,
     /// Read disk IO on this node
@@ -83,22 +83,29 @@ impl TryFrom<HeartbeatRequest> for Stat {
         } = value;
 
         match (header, peer, node_stat) {
-            (Some(header), Some(peer), Some(node_stat)) => Ok(Self {
-                timestamp_millis: time_util::current_time_millis(),
-                cluster_id: header.cluster_id,
-                id: peer.id,
-                addr: peer.addr,
-                is_leader,
-                rcus: node_stat.rcus,
-                wcus: node_stat.wcus,
-                table_num: node_stat.table_num,
-                region_num: node_stat.region_num,
-                cpu_usage: node_stat.cpu_usage,
-                load: node_stat.load,
-                read_io_rate: node_stat.read_io_rate,
-                write_io_rate: node_stat.write_io_rate,
-                region_stats: region_stats.into_iter().map(RegionStat::from).collect(),
-            }),
+            (Some(header), Some(peer), Some(node_stat)) => {
+                let region_num = if node_stat.region_num > 0 {
+                    Some(node_stat.region_num as u64)
+                } else {
+                    None
+                };
+                Ok(Self {
+                    timestamp_millis: time_util::current_time_millis(),
+                    cluster_id: header.cluster_id,
+                    id: peer.id,
+                    addr: peer.addr,
+                    is_leader,
+                    rcus: node_stat.rcus,
+                    wcus: node_stat.wcus,
+                    table_num: node_stat.table_num,
+                    region_num,
+                    cpu_usage: node_stat.cpu_usage,
+                    load: node_stat.load,
+                    read_io_rate: node_stat.read_io_rate,
+                    write_io_rate: node_stat.write_io_rate,
+                    region_stats: region_stats.into_iter().map(RegionStat::from).collect(),
+                })
+            }
             _ => Err(()),
         }
     }
@@ -129,7 +136,7 @@ mod tests {
         let stat = Stat {
             cluster_id: 3,
             id: 101,
-            region_num: 10,
+            region_num: Some(10),
             ..Default::default()
         };
 

src/meta-srv/src/handler/persist_stats_handler.rs

@@ -80,14 +80,17 @@ mod tests {
             kv_store,
             election: None,
             skip_all: Arc::new(AtomicBool::new(false)),
+            catalog: None,
+            schema: None,
+            table: None,
         };
 
         let req = HeartbeatRequest::default();
         let mut acc = HeartbeatAccumulator {
             stats: vec![Stat {
                 cluster_id: 3,
                 id: 101,
-                region_num: 100,
+                region_num: Some(100),
                 ..Default::default()
             }],
             ..Default::default()
@@ -122,6 +125,6 @@ mod tests {
         let val: StatValue = kv.value.clone().try_into().unwrap();
 
         assert_eq!(1, val.stats.len());
-        assert_eq!(100, val.stats[0].region_num);
+        assert_eq!(Some(100), val.stats[0].region_num);
     }
 }

src/meta-srv/src/handler/response_header_handler.rs

@@ -62,6 +62,9 @@ mod tests {
             kv_store,
             election: None,
             skip_all: Arc::new(AtomicBool::new(false)),
+            catalog: None,
+            schema: None,
+            table: None,
         };
 
         let req = HeartbeatRequest {