feat(query): atomic workload group metadata management (#17880)

feat(query): atomic workload metadata management

Author: zhang2014

src/common/exception/src/exception_code.rs

@@ -404,6 +404,11 @@ build_exceptions! {
     UnknownProcedure(3130),
     ProcedureAlreadyExists(3131),
     IllegalProcedureFormat(3132),
+    // Workload
+    InvalidWorkload(3140),
+    AlreadyExistsWorkload(3141),
+    UnknownWorkload(3142),
+    WorkloadOperateConflict(3143),
 }
 
 // Storage errors [3001, 4000].

src/query/management/src/lib.rs

@@ -30,6 +30,7 @@ mod warehouse;
 mod client_session;
 pub mod errors;
 mod procedure;
+mod workload;
 
 pub use client_session::ClientSessionMgr;
 pub use connection::ConnectionMgr;
@@ -55,3 +56,7 @@ pub use warehouse::SystemManagedWarehouse;
 pub use warehouse::WarehouseApi;
 pub use warehouse::WarehouseInfo;
 pub use warehouse::WarehouseMgr;
+pub use workload::QuotaValue;
+pub use workload::WorkloadApi;
+pub use workload::WorkloadGroup;
+pub use workload::WorkloadMgr;

src/query/management/src/workload/mod.rs

@@ -0,0 +1,21 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+mod workload_api;
+mod workload_mgr;
+
+pub use workload_api::QuotaValue;
+pub use workload_api::WorkloadApi;
+pub use workload_api::WorkloadGroup;
+pub use workload_mgr::WorkloadMgr;

src/query/management/src/workload/workload_api.rs

@@ -0,0 +1,48 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::HashMap;
+use std::time::Duration;
+
+use databend_common_exception::Result;
+
+#[derive(serde::Serialize, serde::Deserialize, Clone, Eq, PartialEq, Debug)]
+pub enum QuotaValue {
+    Duration(Duration),
+    String(String),
+}
+
+#[derive(serde::Serialize, serde::Deserialize, Eq, PartialEq, Debug)]
+pub struct WorkloadGroup {
+    pub id: String,
+    pub name: String,
+    pub quotas: HashMap<String, QuotaValue>,
+}
+
+#[async_trait::async_trait]
+pub trait WorkloadApi: Sync + Send {
+    async fn create(&self, group: WorkloadGroup) -> Result<WorkloadGroup>;
+
+    async fn drop(&self, name: String) -> Result<()>;
+
+    async fn rename(&self, old_name: String, new_name: String) -> Result<()>;
+
+    async fn alter_quotas(&self, name: String, quotas: HashMap<String, QuotaValue>) -> Result<()>;
+
+    async fn get_all(&self) -> Result<Vec<WorkloadGroup>>;
+
+    async fn get_by_id(&self, id: &str) -> Result<WorkloadGroup>;
+
+    async fn get_by_name(&self, name: &str) -> Result<WorkloadGroup>;
+}

src/query/management/src/workload/workload_mgr.rs

@@ -0,0 +1,263 @@
+// Copyright 2021 Datafuse Labs
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+use std::collections::HashMap;
+
+use databend_common_base::base::escape_for_key;
+use databend_common_base::base::GlobalUniqName;
+use databend_common_exception::ErrorCode;
+use databend_common_exception::Result;
+use databend_common_meta_kvapi::kvapi::KVApi;
+use databend_common_meta_store::MetaStore;
+use databend_common_meta_types::SeqV;
+use databend_common_meta_types::SeqValue;
+use databend_common_meta_types::TxnCondition;
+use databend_common_meta_types::TxnOp;
+use databend_common_meta_types::TxnRequest;
+
+use crate::workload::workload_api::QuotaValue;
+use crate::workload::workload_api::WorkloadApi;
+use crate::workload::workload_api::WorkloadGroup;
+pub static WORKLOAD_META_KEY_PREFIX: &str = "__fd_workloads";
+
+pub struct WorkloadMgr {
+    metastore: MetaStore,
+    workload_key_prefix: String,
+    workload_index_prefix: String,
+}
+
+impl WorkloadMgr {
+    pub fn create(metastore: MetaStore, tenant: &str) -> Result<Self> {
+        Ok(WorkloadMgr {
+            metastore,
+            workload_key_prefix: format!(
+                "{}/{}/workload_groups",
+                WORKLOAD_META_KEY_PREFIX,
+                escape_for_key(tenant)?
+            ),
+            workload_index_prefix: format!(
+                "{}/{}/index_workload_groups",
+                WORKLOAD_META_KEY_PREFIX,
+                escape_for_key(tenant)?
+            ),
+        })
+    }
+
+    pub async fn get_id_by_name(&self, name: &str) -> Result<String> {
+        let index_key = format!("{}/{}", self.workload_index_prefix, escape_for_key(name)?);
+
+        let Some(seq) = self.metastore.get_kv(&index_key).await? else {
+            return Err(ErrorCode::UnknownWorkload(format!(
+                "Unknown workload {}",
+                name
+            )));
+        };
+
+        Ok(unsafe { String::from_utf8_unchecked(seq.data) })
+    }
+
+    async fn get_seq_by_name(&self, name: &str) -> Result<SeqV<WorkloadGroup>> {
+        let id = self.get_id_by_name(name).await?;
+
+        let Some(seq) = self.get_seq_by_id(&id).await? else {
+            return Err(ErrorCode::UnknownWorkload(format!(
+                "Unknown workload {}",
+                name
+            )));
+        };
+
+        Ok(seq)
+    }
+
+    async fn get_seq_by_id(&self, id: &str) -> Result<Option<SeqV<WorkloadGroup>>> {
+        let workload_key = format!("{}/{}", self.workload_key_prefix, id);
+
+        let Some(seq) = self.metastore.get_kv(&workload_key).await? else {
+            return Ok(None);
+        };
+
+        Ok(Some(SeqV::new(
+            seq.seq,
+            serde_json::from_slice::<WorkloadGroup>(&seq.data)?,
+        )))
+    }
+}
+
+#[async_trait::async_trait]
+impl WorkloadApi for WorkloadMgr {
+    async fn create(&self, mut group: WorkloadGroup) -> Result<WorkloadGroup> {
+        if group.name.is_empty() {
+            return Err(ErrorCode::InvalidWorkload("Workload group name is empty."));
+        }
+
+        if !group.id.is_empty() {
+            return Err(ErrorCode::InvalidWorkload(
+                "Workload group id is not empty.",
+            ));
+        }
+
+        let group_id = GlobalUniqName::unique();
+        let mut create_workload = TxnRequest::default();
+
+        group.id = group_id.clone();
+        let escape_name = escape_for_key(&group.name)?;
+        let workload_key = format!("{}/{}", self.workload_key_prefix, group_id);
+        let workload_index_key = format!("{}/{}", self.workload_index_prefix, escape_name);
+
+        create_workload
+            .condition
+            .push(TxnCondition::eq_seq(workload_index_key.clone(), 0));
+        create_workload
+            .if_then
+            .push(TxnOp::put(workload_index_key, group_id.into_bytes()));
+        create_workload
+            .if_then
+            .push(TxnOp::put(workload_key, serde_json::to_vec(&group)?));
+
+        match self.metastore.transaction(create_workload).await? {
+            res if res.success => Ok(group),
+            _res => Err(ErrorCode::AlreadyExistsWorkload(format!(
+                "The workload {} already exits.",
+                group.name
+            ))),
+        }
+    }
+
+    async fn drop(&self, name: String) -> Result<()> {
+        let workload_id = self.get_id_by_name(&name).await?;
+
+        let escape_name = escape_for_key(&name)?;
+        let workload_key = format!("{}/{}", self.workload_key_prefix, workload_id);
+        let workload_index_key = format!("{}/{}", self.workload_index_prefix, escape_name);
+
+        let mut drop_workload = TxnRequest::default();
+        drop_workload.condition.push(TxnCondition::eq_value(
+            workload_index_key.clone(),
+            workload_id.into_bytes(),
+        ));
+        drop_workload.if_then.push(TxnOp::delete(workload_key));
+        drop_workload
+            .if_then
+            .push(TxnOp::delete(workload_index_key));
+
+        match self.metastore.transaction(drop_workload).await? {
+            res if res.success => Ok(()),
+            _res => Err(ErrorCode::UnknownWorkload(format!(
+                "Unknown workload {}",
+                name
+            ))),
+        }
+    }
+
+    async fn rename(&self, old_name: String, new_name: String) -> Result<()> {
+        let mut workload = self.get_by_name(&old_name).await?;
+        workload.name = new_name.clone();
+
+        let workload_key = format!("{}/{}", self.workload_key_prefix, workload.id);
+
+        let escape_name = escape_for_key(&old_name)?;
+        let old_workload_index_key = format!("{}/{}", self.workload_index_prefix, escape_name);
+
+        let escape_name = escape_for_key(&new_name)?;
+        let new_workload_index_key = format!("{}/{}", self.workload_index_prefix, escape_name);
+
+        let mut rename_workload = TxnRequest::default();
+        rename_workload
+            .condition
+            .push(TxnCondition::eq_seq(new_workload_index_key.clone(), 0));
+        rename_workload.condition.push(TxnCondition::eq_value(
+            old_workload_index_key.clone(),
+            workload.id.clone().into_bytes(),
+        ));
+
+        rename_workload
+            .if_then
+            .push(TxnOp::put(workload_key, serde_json::to_vec(&workload)?));
+        rename_workload
+            .if_then
+            .push(TxnOp::delete(old_workload_index_key));
+        rename_workload.if_then.push(TxnOp::put(
+            new_workload_index_key,
+            workload.id.clone().into_bytes(),
+        ));
+
+        match self.metastore.transaction(rename_workload).await? {
+            res if res.success => Ok(()),
+            _res => Err(ErrorCode::InvalidWorkload(format!(
+                "Unknown workload {} or workload {} already exists",
+                old_name, new_name
+            ))),
+        }
+    }
+
+    async fn alter_quotas(&self, name: String, quotas: HashMap<String, QuotaValue>) -> Result<()> {
+        for _index in 0..5 {
+            let workload = self.get_seq_by_name(&name).await?;
+            let seq = workload.seq;
+            let mut workload = workload.into_value().unwrap();
+
+            for (key, value) in &quotas {
+                workload.quotas.insert(key.clone(), value.clone());
+            }
+
+            let workload_key = format!("{}/{}", self.workload_key_prefix, workload.id);
+            let mut alter_workload = TxnRequest::default();
+            alter_workload
+                .condition
+                .push(TxnCondition::eq_seq(workload_key.clone(), seq));
+            alter_workload
+                .if_then
+                .push(TxnOp::put(workload_key, serde_json::to_vec(&workload)?));
+
+            if self.metastore.transaction(alter_workload).await?.success {
+                return Ok(());
+            }
+        }
+
+        Err(ErrorCode::WorkloadOperateConflict(
+            "Workload operate conflict(tried 5 times).",
+        ))
+    }
+
+    async fn get_all(&self) -> Result<Vec<WorkloadGroup>> {
+        let list_reply = self
+            .metastore
+            .prefix_list_kv(&format!("{}/", self.workload_key_prefix))
+            .await?;
+
+        let mut workload_groups = Vec::with_capacity(list_reply.len());
+
+        for (_key, seq) in list_reply {
+            workload_groups.push(serde_json::from_slice::<WorkloadGroup>(&seq.data)?);
+        }
+
+        Ok(workload_groups)
+    }
+
+    async fn get_by_id(&self, id: &str) -> Result<WorkloadGroup> {
+        let Some(seq) = self.get_seq_by_id(id).await? else {
+            return Err(ErrorCode::UnknownWorkload(format!(
+                "Unknown workload id {}",
+                id
+            )));
+        };
+
+        Ok(seq.into_value().unwrap())
+    }
+
+    async fn get_by_name(&self, name: &str) -> Result<WorkloadGroup> {
+        let seq_value = self.get_seq_by_name(name).await?;
+        Ok(seq_value.into_value().unwrap())
+    }
+}

src/query/management/tests/it/main.rs

@@ -22,3 +22,4 @@ mod stage;
 mod udf;
 mod user;
 mod warehouse;
+mod workload;