chore(query): add max node quota (#18474)

* chore(query): add max node quota

* chore(query): add max node quota

* chore(query): add max node quota

* chore(query): add max node quota

Author: zhang2014

Cargo.lock

@@ -22,8 +22,8 @@ use serde::Serialize;
 
 #[derive(Debug, Clone, Eq, Ord, PartialOrd, PartialEq, serde::Serialize, serde::Deserialize)]
 pub struct ComputeQuota {
-    threads_num: Option<usize>,
-    memory_usage: Option<usize>,
+    pub threads_num: Option<usize>,
+    pub memory_usage: Option<usize>,
 }
 
 #[derive(Debug, Clone, Eq, Ord, PartialOrd, PartialEq, serde::Serialize, serde::Deserialize)]
@@ -85,6 +85,10 @@ pub enum Feature {
     VectorIndex,
     #[serde(alias = "private_task", alias = "PRIVATE_TASK")]
     PrivateTask,
+    #[serde(alias = "max_node_quota", alias = "MAX_NODE_QUOTA")]
+    MaxNodeQuota(usize),
+    #[serde(alias = "max_cpu_quota", alias = "MAX_CPU_QUOTA")]
+    MaxCpuQuota(usize),
     #[serde(other)]
     Unknown,
 }
@@ -138,13 +142,18 @@ impl fmt::Display for Feature {
             Feature::VectorIndex => write!(f, "vector_index"),
             Feature::PrivateTask => write!(f, "private_task"),
             Feature::Unknown => write!(f, "unknown"),
+            Feature::MaxCpuQuota(v) => write!(f, "max_cpu_quota({})", v),
+            Feature::MaxNodeQuota(v) => write!(f, "max_node_quota({})", v),
         }
     }
 }
 
 impl Feature {
     pub fn verify_default(&self, message: impl Into<String>) -> Result<(), ErrorCode> {
-        Err(ErrorCode::LicenseKeyInvalid(message.into()))
+        match self {
+            Feature::MaxCpuQuota(_) | Feature::MaxNodeQuota(_) => Ok(()),
+            _ => Err(ErrorCode::LicenseKeyInvalid(message.into())),
+        }
     }
 
     pub fn verify(&self, feature: &Feature) -> Result<bool, ErrorCode> {
@@ -173,6 +182,8 @@ impl Feature {
 
                 Ok(true)
             }
+            (Feature::MaxCpuQuota(c), Feature::MaxCpuQuota(v)) => Ok(c > v),
+            (Feature::MaxNodeQuota(c), Feature::MaxNodeQuota(v)) => Ok(c > v),
             (Feature::Test, Feature::Test)
             | (Feature::AggregateIndex, Feature::AggregateIndex)
             | (Feature::ComputedColumn, Feature::ComputedColumn)
@@ -380,6 +391,11 @@ mod tests {
             serde_json::from_str::<Feature>("\"private_task\"").unwrap()
         );
 
+        assert_eq!(
+            Feature::MaxNodeQuota(1),
+            serde_json::from_str::<Feature>("{\"MaxNodeQuota\": 1}").unwrap()
+        );
+
         assert_eq!(
             Feature::Unknown,
             serde_json::from_str::<Feature>("\"ssss\"").unwrap()

src/common/license/src/license.rs

@@ -19,6 +19,7 @@ log = { workspace = true }
 map-api = { workspace = true }
 num-derive = { workspace = true }
 num-traits = { workspace = true }
+num_cpus = { workspace = true }
 openraft = { workspace = true }
 pretty_assertions = { workspace = true }
 prost = { workspace = true }

src/meta/types/Cargo.toml

@@ -56,13 +56,13 @@ impl NodeInfo {
     pub fn create(
         id: String,
         secret: String,
-        cpu_nums: u64,
         http_address: String,
         flight_address: String,
         discovery_address: String,
         binary_version: String,
         cache_id: String,
     ) -> NodeInfo {
+        let cpu_nums = num_cpus::get() as u64;
         NodeInfo {
             id,
             secret,

src/meta/types/src/cluster.rs

@@ -79,7 +79,6 @@ fn create_node(cache_id: String) -> Arc<NodeInfo> {
     Arc::new(NodeInfo::create(
         GlobalUniqName::unique(),
         String::new(),
-        0,
         String::new(),
         String::new(),
         String::new(),

src/query/catalog/tests/it/partitions.rs

@@ -43,13 +43,17 @@ use databend_common_config::InnerConfig;
 use databend_common_exception::ErrorCode;
 use databend_common_exception::Result;
 use databend_common_grpc::ConnectionFactory;
+use databend_common_license::license::Feature;
+use databend_common_license::license_manager::LicenseManagerSwitch;
 use databend_common_management::WarehouseApi;
 use databend_common_management::WarehouseMgr;
+use databend_common_meta_app::tenant::Tenant;
 use databend_common_meta_store::MetaStore;
 use databend_common_meta_store::MetaStoreProvider;
 use databend_common_meta_types::NodeInfo;
 use databend_common_meta_types::SeqV;
 use databend_common_metrics::cluster::*;
+use databend_common_settings::Settings;
 use databend_common_version::DATABEND_COMMIT_VERSION;
 use databend_enterprise_resources_management::ResourcesManagement;
 use futures::future::select;
@@ -322,7 +326,6 @@ impl ClusterDiscovery {
                     let mut node_info = NodeInfo::create(
                         config.query.node_id.clone(),
                         config.query.node_secret.clone(),
-                        config.query.num_cpus,
                         format!(
                             "{}:{}",
                             config.query.http_handler_host, config.query.http_handler_port
@@ -387,7 +390,6 @@ impl ClusterDiscovery {
                 let mut node_info = NodeInfo::create(
                     config.query.node_id.clone(),
                     config.query.node_secret.clone(),
-                    config.query.num_cpus,
                     format!(
                         "{}:{}",
                         config.query.http_handler_host, config.query.http_handler_port
@@ -535,7 +537,6 @@ impl ClusterDiscovery {
 
     #[async_backtrace::framed]
     pub async fn register_to_metastore(self: &Arc<Self>, cfg: &InnerConfig) -> Result<()> {
-        let cpus = cfg.query.num_cpus;
         let mut address = cfg.query.flight_api_address.clone();
         let mut http_address = format!(
             "{}:{}",
@@ -601,7 +602,6 @@ impl ClusterDiscovery {
         let mut node_info = NodeInfo::create(
             self.local_id.clone(),
             self.local_secret.clone(),
-            cpus,
             http_address,
             address,
             discovery_address,
@@ -614,6 +614,9 @@ impl ClusterDiscovery {
 
         self.drop_invalid_nodes(&node_info).await?;
 
+        let online_nodes = self.warehouse_manager.list_online_nodes().await?;
+        self.check_license_key(online_nodes).await?;
+
         match self.warehouse_manager.start_node(node_info).await {
             Ok(seq_node) => self.start_heartbeat(seq_node).await,
             Err(cause) => Err(cause.add_message_back("(while cluster api add_node).")),
@@ -628,6 +631,24 @@ impl ClusterDiscovery {
         heartbeat.start(node_info, seq);
         Ok(())
     }
+
+    async fn check_license_key(&self, nodes: Vec<NodeInfo>) -> Result<()> {
+        let license_key = Self::get_license_key(&self.tenant_id).await?;
+
+        let total_cpu_nums = nodes.iter().map(|x| x.cpu_nums).sum::<u64>();
+        LicenseManagerSwitch::instance()
+            .check_enterprise_enabled(license_key.clone(), Feature::MaxNodeQuota(nodes.len()))?;
+
+        LicenseManagerSwitch::instance()
+            .check_enterprise_enabled(license_key, Feature::MaxCpuQuota(total_cpu_nums as usize))
+    }
+
+    async fn get_license_key(tenant: &str) -> Result<String> {
+        // We must get the license key from settings. It may be in the configuration file.
+        let settings = Settings::create(Tenant::new_literal(tenant));
+        settings.load_changes().await?;
+        Ok(settings.get_enterprise_license())
+    }
 }
 
 struct ClusterHeartbeat {

src/query/service/src/clusters/cluster.rs

@@ -38,7 +38,6 @@ impl ClusterDescriptor {
         new_nodes.push(Arc::new(NodeInfo::create(
             id.clone(),
             "".to_string(),
-            0,
             "".to_string(),
             addr.into(),
             "".to_string(),

src/query/service/src/test_kits/cluster.rs

@@ -507,7 +507,6 @@ fn create_node(local_id: &str) -> Arc<NodeInfo> {
     let mut node_info = NodeInfo::create(
         local_id.to_string(),
         String::new(),
-        0,
         String::new(),
         String::new(),
         String::new(),