feat: add table lock manager to fix table lock cannot be released

src/query/catalog/src/catalog/interface.rs

@@ -75,6 +75,7 @@ use common_meta_app::schema::UpsertTableOptionReply;
 use common_meta_app::schema::UpsertTableOptionReq;
 use common_meta_app::schema::VirtualColumnMeta;
 use common_meta_types::MetaId;
+use common_pipeline_core::TableLockReq;
 use dyn_clone::DynClone;
 
 use crate::database::Database;
@@ -259,22 +260,16 @@ pub trait Catalog: DynClone + Send + Sync + Debug {
         req: TruncateTableReq,
     ) -> Result<TruncateTableReply>;
 
-    async fn list_table_lock_revs(&self, table_id: u64) -> Result<Vec<u64>>;
+    async fn list_table_lock_revs(&self, req: Box<dyn TableLockReq>) -> Result<Vec<u64>>;
 
     async fn create_table_lock_rev(
         &self,
-        expire_secs: u64,
-        table_info: &TableInfo,
+        req: Box<dyn TableLockReq>,
     ) -> Result<CreateTableLockRevReply>;
 
-    async fn extend_table_lock_rev(
-        &self,
-        expire_secs: u64,
-        table_info: &TableInfo,
-        revision: u64,
-    ) -> Result<()>;
+    async fn extend_table_lock_rev(&self, req: Box<dyn TableLockReq>) -> Result<()>;
 
-    async fn delete_table_lock_rev(&self, table_info: &TableInfo, revision: u64) -> Result<()>;
+    async fn delete_table_lock_rev(&self, req: Box<dyn TableLockReq>) -> Result<()>;
 
     /// Table function
 

src/query/ee-features/table-lock/Cargo.toml

@@ -18,12 +18,12 @@ common-catalog = { path = "../../catalog" }
 common-exception = { path = "../../../common/exception" }
 common-license = { path = "../../../common/license" }
 common-meta-app = { path = "../../../meta/app" }
+common-meta-kvapi = { path = "../../../meta/kvapi" }
+common-pipeline-core = { path = "../../pipeline/core" }
 
 async-backtrace = { workspace = true }
 async-trait = "0.1.57"
-futures = "0.3.24"
-log = { workspace = true }
-rand = "0.8.5"
+chrono = { workspace = true }
 
 [build-dependencies]
 

src/query/ee-features/table-lock/src/lib.rs

@@ -12,9 +12,13 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-mod table_lock_handler;
-mod table_lock_heartbeat;
+mod table_level_lock;
+mod table_lock_manager;
 
-pub use table_lock_handler::TableLockHandler;
-pub use table_lock_handler::TableLockHandlerWrapper;
-pub use table_lock_heartbeat::TableLockHeartbeat;
+pub use table_level_lock::CreateTableLockReq;
+pub use table_level_lock::DeleteTableLockReq;
+pub use table_level_lock::ExtendTableLockReq;
+pub use table_level_lock::ListTableLockReq;
+pub use table_level_lock::TableLevelLock;
+pub use table_lock_manager::TableLockManager;
+pub use table_lock_manager::TableLockManagerWrapper;

src/query/ee-features/table-lock/src/table_level_lock/lock.rs

@@ -0,0 +1,104 @@
+// 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::sync::Arc;
+
+use common_meta_app::schema::TableLockKey;
+use common_meta_kvapi::kvapi::Key;
+use common_pipeline_core::TableLock;
+use common_pipeline_core::TableLockLevel;
+use common_pipeline_core::TableLockReq;
+
+use crate::table_level_lock::req::*;
+use crate::TableLockManagerWrapper;
+
+pub struct TableLevelLock {
+    lock_mgr: Arc<TableLockManagerWrapper>,
+    table_id: u64,
+    revision: u64,
+}
+
+impl TableLevelLock {
+    pub fn create(lock_mgr: Arc<TableLockManagerWrapper>, table_id: u64) -> Self {
+        TableLevelLock {
+            lock_mgr,
+            table_id,
+            revision: 0,
+        }
+    }
+}
+
+#[async_trait::async_trait]
+impl TableLock for TableLevelLock {
+    fn level(&self) -> TableLockLevel {
+        TableLockLevel::Table
+    }
+
+    fn table_id(&self) -> u64 {
+        self.table_id
+    }
+
+    fn set_revision(&mut self, revision: u64) {
+        self.revision = revision;
+    }
+
+    fn revision(&self) -> u64 {
+        self.revision
+    }
+
+    fn watch_delete_key(&self, revision: u64) -> String {
+        let lock_key = TableLockKey {
+            table_id: self.table_id,
+            revision,
+        };
+        lock_key.to_string_key()
+    }
+
+    fn create_table_lock_req(&self, expire_secs: u64) -> Box<dyn TableLockReq> {
+        Box::new(CreateTableLockReq {
+            table_id: self.table_id,
+            expire_secs,
+        })
+    }
+
+    fn extend_table_lock_req(&self, expire_secs: u64) -> Box<dyn TableLockReq> {
+        Box::new(ExtendTableLockReq {
+            table_id: self.table_id,
+            expire_secs,
+            revision: self.revision,
+        })
+    }
+
+    fn delete_table_lock_req(&self) -> Box<dyn TableLockReq> {
+        Box::new(DeleteTableLockReq {
+            table_id: self.table_id,
+            revision: self.revision,
+        })
+    }
+
+    fn list_table_lock_req(&self) -> Box<dyn TableLockReq> {
+        Box::new(ListTableLockReq {
+            table_id: self.table_id,
+        })
+    }
+}
+
+impl Drop for TableLevelLock {
+    fn drop(&mut self) {
+        let revision = self.revision();
+        if revision > 0 {
+            self.lock_mgr.unlock(revision);
+        }
+    }
+}

src/query/ee-features/table-lock/src/table_level_lock/mod.rs

@@ -0,0 +1,22 @@
+// 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 lock;
+mod req;
+
+pub use lock::TableLevelLock;
+pub use req::CreateTableLockReq;
+pub use req::DeleteTableLockReq;
+pub use req::ExtendTableLockReq;
+pub use req::ListTableLockReq;

src/query/ee-features/table-lock/src/table_level_lock/req.rs

@@ -0,0 +1,122 @@
+// 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::any::Any;
+
+use chrono::Utc;
+use common_meta_app::schema::CreateTableLockRevReq;
+use common_meta_app::schema::DeleteTableLockRevReq;
+use common_meta_app::schema::ExtendTableLockRevReq;
+use common_meta_app::schema::ListTableLockRevReq;
+use common_pipeline_core::TableLockReq;
+
+#[derive(Clone)]
+pub struct ListTableLockReq {
+    pub table_id: u64,
+}
+
+impl TableLockReq for ListTableLockReq {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn clone_self(&self) -> Box<dyn TableLockReq> {
+        Box::new(self.clone())
+    }
+}
+
+impl From<&ListTableLockReq> for ListTableLockRevReq {
+    fn from(value: &ListTableLockReq) -> Self {
+        ListTableLockRevReq {
+            table_id: value.table_id,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct CreateTableLockReq {
+    pub table_id: u64,
+    pub expire_secs: u64,
+}
+
+impl TableLockReq for CreateTableLockReq {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn clone_self(&self) -> Box<dyn TableLockReq> {
+        Box::new(self.clone())
+    }
+}
+
+impl From<&CreateTableLockReq> for CreateTableLockRevReq {
+    fn from(value: &CreateTableLockReq) -> Self {
+        CreateTableLockRevReq {
+            table_id: value.table_id,
+            expire_at: Utc::now().timestamp() as u64 + value.expire_secs,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct ExtendTableLockReq {
+    pub table_id: u64,
+    pub expire_secs: u64,
+    pub revision: u64,
+}
+
+impl TableLockReq for ExtendTableLockReq {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn clone_self(&self) -> Box<dyn TableLockReq> {
+        Box::new(self.clone())
+    }
+}
+
+impl From<&ExtendTableLockReq> for ExtendTableLockRevReq {
+    fn from(value: &ExtendTableLockReq) -> Self {
+        ExtendTableLockRevReq {
+            table_id: value.table_id,
+            expire_at: Utc::now().timestamp() as u64 + value.expire_secs,
+            revision: value.revision,
+        }
+    }
+}
+
+#[derive(Clone)]
+pub struct DeleteTableLockReq {
+    pub table_id: u64,
+    pub revision: u64,
+}
+
+impl TableLockReq for DeleteTableLockReq {
+    fn as_any(&self) -> &dyn Any {
+        self
+    }
+
+    fn clone_self(&self) -> Box<dyn TableLockReq> {
+        Box::new(self.clone())
+    }
+}
+
+impl From<&DeleteTableLockReq> for DeleteTableLockRevReq {
+    fn from(value: &DeleteTableLockReq) -> Self {
+        DeleteTableLockRevReq {
+            table_id: value.table_id,
+            revision: value.revision,
+        }
+    }
+}