fix: task state will be affected by get_next_ready_tasks

Author: KKould

src/meta/app/src/principal/mod.rs

@@ -101,7 +101,7 @@ pub use task::TaskDependentKey;
 pub use task::TaskDependentValue;
 pub use task::TaskMessage;
 pub use task::TaskRun;
-pub use task::TaskState;
+pub use task::TaskStateValue;
 pub use task::WarehouseOptions;
 pub use task_ident::TaskIdent;
 pub use task_ident::TaskIdentRaw;

src/meta/app/src/principal/task.rs

@@ -175,7 +175,19 @@ impl TaskMessage {
 }
 
 #[derive(Debug, Clone, PartialEq)]
-pub struct TaskState {
+pub struct TaskStateKey {
+    pub current: String,
+    pub next: String,
+}
+
+impl TaskStateKey {
+    pub fn new(current: String, next: String) -> Self {
+        Self { current, next }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub struct TaskStateValue {
     pub is_succeeded: bool,
 }
 
@@ -202,11 +214,29 @@ pub struct TaskDependentValue(pub BTreeSet<String>);
 
 mod kvapi_key_impl {
     use databend_common_meta_kvapi::kvapi;
+    use databend_common_meta_kvapi::kvapi::KeyBuilder;
     use databend_common_meta_kvapi::kvapi::KeyError;
+    use databend_common_meta_kvapi::kvapi::KeyParser;
 
+    use crate::principal::task::TaskStateKey;
     use crate::principal::DependentType;
     use crate::principal::TaskDependentKey;
 
+    impl kvapi::KeyCodec for TaskStateKey {
+        fn encode_key(&self, b: KeyBuilder) -> KeyBuilder {
+            b.push_str(self.current.as_str())
+                .push_str(self.next.as_str())
+        }
+
+        fn decode_key(parser: &mut KeyParser) -> Result<Self, KeyError>
+        where Self: Sized {
+            let current = parser.next_str()?;
+            let next = parser.next_str()?;
+
+            Ok(Self { current, next })
+        }
+    }
+
     impl kvapi::KeyCodec for TaskDependentKey {
         fn encode_key(&self, b: kvapi::KeyBuilder) -> kvapi::KeyBuilder {
             match self.ty {

src/meta/app/src/principal/task_state_ident.rs

@@ -14,16 +14,26 @@
 
 use crate::tenant_key::ident::TIdent;
 
-pub type TaskStateIdent = TIdent<Resource>;
-
-pub type TaskStateIdentRaw = TIdent<Resource>;
+pub type TaskStateIdent = TIdent<Resource, TaskStateKey>;
+
+impl TaskStateIdent {
+    pub fn new(tenant: impl ToTenant, current: impl ToString, next: impl ToString) -> Self {
+        TaskStateIdent::new_generic(
+            tenant,
+            TaskStateKey::new(current.to_string(), next.to_string()),
+        )
+    }
+}
 
 pub use kvapi_impl::Resource;
 
+use crate::principal::task::TaskStateKey;
+use crate::tenant::ToTenant;
+
 mod kvapi_impl {
     use databend_common_meta_kvapi::kvapi;
 
-    use crate::principal::task::TaskState;
+    use crate::principal::task::TaskStateValue;
     use crate::principal::task_state_ident::TaskStateIdent;
     use crate::tenant_key::resource::TenantResource;
 
@@ -32,10 +42,10 @@ mod kvapi_impl {
         const PREFIX: &'static str = "__fd_task_states";
         const TYPE: &'static str = "TaskStateIdent";
         const HAS_TENANT: bool = true;
-        type ValueType = TaskState;
+        type ValueType = TaskStateValue;
     }
 
-    impl kvapi::Value for TaskState {
+    impl kvapi::Value for TaskStateValue {
         type KeyType = TaskStateIdent;
 
         fn dependency_keys(&self, _key: &Self::KeyType) -> impl IntoIterator<Item = String> {

src/meta/proto-conv/src/task_from_to_protobuf_impl.rs

@@ -219,8 +219,8 @@ impl FromToProto for mt::TaskDependentValue {
     }
 }
 
-impl FromToProto for mt::TaskState {
-    type PB = pb::TaskState;
+impl FromToProto for mt::TaskStateValue {
+    type PB = pb::TaskStateValue;
 
     fn get_pb_ver(p: &Self::PB) -> u64 {
         p.ver
@@ -234,7 +234,7 @@ impl FromToProto for mt::TaskState {
     }
 
     fn to_pb(&self) -> Result<Self::PB, Incompatible> {
-        Ok(pb::TaskState {
+        Ok(pb::TaskStateValue {
             ver: VER,
             min_reader_ver: MIN_READER_VER,
             is_succeeded: self.is_succeeded,

src/meta/proto-conv/tests/it/v141_task_state.rs

@@ -23,7 +23,7 @@ use crate::common;
 fn test_decode_v141_task_state() -> anyhow::Result<()> {
     let task_state_v141 = vec![8, 1, 160, 6, 141, 1, 168, 6, 24];
 
-    let want = || mt::TaskState { is_succeeded: true };
+    let want = || mt::TaskStateValue { is_succeeded: true };
     common::test_pb_from_to(func_name!(), want())?;
     common::test_load_old(func_name!(), task_state_v141.as_slice(), 141, want())?;
 

src/meta/protos/proto/task.proto

@@ -79,25 +79,13 @@ message Task {
   string owner_user = 21;
 }
 
-message TaskState {
+message TaskStateValue {
   uint64 ver = 100;
   uint64 min_reader_ver = 101;
 
   bool is_succeeded = 1;
 }
 
-message TaskDependentKey {
-  uint64 ver = 100;
-  uint64 min_reader_ver = 101;
-
-  enum DependentType {
-    After = 0;
-    Before = 1;
-  }
-  DependentType ty = 1;
-  string source = 2;
-}
-
 message TaskDependentValue {
   uint64 ver = 100;
   uint64 min_reader_ver = 101;

src/query/management/src/task/task_mgr.rs

@@ -27,7 +27,7 @@ use databend_common_meta_api::util::txn_put_pb;
 use databend_common_meta_api::SchemaApi;
 use databend_common_meta_app::principal::task;
 use databend_common_meta_app::principal::task::TaskMessage;
-use databend_common_meta_app::principal::task::TaskState;
+use databend_common_meta_app::principal::task::TaskStateValue;
 use databend_common_meta_app::principal::task_dependent_ident::TaskDependentIdent;
 use databend_common_meta_app::principal::task_message_ident::TaskMessageIdent;
 use databend_common_meta_app::principal::task_state_ident::TaskStateIdent;
@@ -47,6 +47,7 @@ use databend_common_meta_types::TxnOp;
 use databend_common_meta_types::TxnRequest;
 use databend_common_meta_types::With;
 use databend_common_proto_conv::FromToProto;
+use futures::StreamExt;
 use futures::TryStreamExt;
 use prost::Message;
 use seq_marked::SeqValue;
@@ -266,37 +267,38 @@ impl TaskMgr {
         task_name: &str,
         new_afters: &[String],
     ) -> Result<Result<(), TaskError>, TaskApiError> {
-        let mut update_ops = Vec::new();
-        let mut check_ops = Vec::with_capacity(new_afters.len());
-
+        let mut request = TxnRequest::new(vec![], vec![]);
         let after_dependent_ident = TaskDependentIdent::new_after(&self.tenant, task_name);
 
         let after_seq_deps = self.kv_api.get_pb(&after_dependent_ident).await?;
-        check_ops.push(txn_cond_eq_seq(
+        request.condition.push(txn_cond_eq_seq(
             &after_dependent_ident,
             after_seq_deps.seq(),
         ));
 
         let mut after_deps = after_seq_deps.unwrap_or_default();
         after_deps.0.extend(new_afters.iter().cloned());
-        update_ops.push(txn_put_pb(&after_dependent_ident, &after_deps)?);
+        request
+            .if_then
+            .push(txn_put_pb(&after_dependent_ident, &after_deps)?);
 
         let before_dependent_idents = new_afters
             .iter()
             .map(|after| TaskDependentIdent::new_before(&self.tenant, after));
         for (before_dependent_ident, before_seq_deps) in
             self.kv_api.get_pb_vec(before_dependent_idents).await?
         {
-            check_ops.push(txn_cond_eq_seq(
+            request.condition.push(txn_cond_eq_seq(
                 &before_dependent_ident,
                 before_seq_deps.seq(),
             ));
 
             let mut deps = before_seq_deps.unwrap_or_default();
             deps.0.insert(task_name.to_string());
-            update_ops.push(txn_put_pb(&before_dependent_ident, &deps)?);
+            request
+                .if_then
+                .push(txn_put_pb(&before_dependent_ident, &deps)?);
         }
-        let request = TxnRequest::new(check_ops, update_ops);
         let reply = self.kv_api.transaction(request).await?;
 
         if !reply.success {
@@ -316,12 +318,11 @@ impl TaskMgr {
         task_name: &str,
         remove_afters: &[String],
     ) -> Result<Result<(), TaskError>, TaskApiError> {
-        let mut update_ops = Vec::new();
-        let mut check_ops = Vec::with_capacity(remove_afters.len());
+        let mut request = TxnRequest::new(vec![], vec![]);
 
         let after_dependent_ident = TaskDependentIdent::new_after(&self.tenant, task_name);
         let after_seq_deps = self.kv_api.get_pb(&after_dependent_ident).await?;
-        check_ops.push(txn_cond_eq_seq(
+        request.condition.push(txn_cond_eq_seq(
             &after_dependent_ident,
             after_seq_deps.seq(),
         ));
@@ -330,7 +331,9 @@ impl TaskMgr {
             for remove_after in remove_afters {
                 deps.0.remove(remove_after);
             }
-            update_ops.push(txn_put_pb(&after_dependent_ident, &deps)?);
+            request
+                .if_then
+                .push(txn_put_pb(&after_dependent_ident, &deps)?);
         }
 
         let before_dependent_idents = remove_afters
@@ -339,17 +342,18 @@ impl TaskMgr {
         for (before_dependent_ident, before_seq_deps) in
             self.kv_api.get_pb_vec(before_dependent_idents).await?
         {
-            check_ops.push(txn_cond_eq_seq(
+            request.condition.push(txn_cond_eq_seq(
                 &before_dependent_ident,
                 before_seq_deps.seq(),
             ));
 
             if let Some(mut deps) = before_seq_deps {
                 deps.0.remove(task_name);
-                update_ops.push(txn_put_pb(&before_dependent_ident, &deps)?);
+                request
+                    .if_then
+                    .push(txn_put_pb(&before_dependent_ident, &deps)?);
             }
         }
-        let request = TxnRequest::new(check_ops, update_ops);
         let reply = self.kv_api.transaction(request).await?;
 
         if !reply.success {
@@ -369,8 +373,7 @@ impl TaskMgr {
         &self,
         task_name: &str,
     ) -> Result<Result<(), TaskError>, TaskApiError> {
-        let mut check_ops = Vec::new();
-        let mut update_ops = Vec::new();
+        let mut request = TxnRequest::new(vec![], vec![]);
 
         let task_after_ident = TaskDependentIdent::new_after(&self.tenant, task_name);
         let task_before_ident = TaskDependentIdent::new_before(&self.tenant, task_name);
@@ -387,24 +390,33 @@ impl TaskMgr {
             }));
         }
         for (target_ident, seq_dep) in self.kv_api.get_pb_vec(target_idents).await? {
-            check_ops.push(txn_cond_eq_seq(&target_ident, seq_dep.seq()));
+            request
+                .condition
+                .push(txn_cond_eq_seq(&target_ident, seq_dep.seq()));
 
             if let Some(mut deps) = seq_dep {
                 deps.0.remove(task_name);
-                update_ops.push(txn_put_pb(&target_ident, &deps)?);
+                request.if_then.push(txn_put_pb(&target_ident, &deps)?);
             }
         }
-        update_ops.push(TxnOp::delete(
+        request.if_then.push(TxnOp::delete(
             TaskDependentIdent::new_before(&self.tenant, task_name).to_string_key(),
         ));
-        update_ops.push(TxnOp::delete(
+        request.if_then.push(TxnOp::delete(
             TaskDependentIdent::new_after(&self.tenant, task_name).to_string_key(),
         ));
-        update_ops.push(TxnOp::delete(
-            TaskStateIdent::new(&self.tenant, task_name).to_string_key(),
-        ));
+        let mut stream = self
+            .kv_api
+            .list_pb_keys(&DirName::new(TaskStateIdent::new(
+                &self.tenant,
+                task_name,
+                "",
+            )))
+            .await?;
 
-        let request = TxnRequest::new(check_ops, update_ops);
+        while let Some(result) = stream.next().await {
+            request.if_then.push(TxnOp::delete(result?.to_string_key()))
+        }
         let _ = self.kv_api.transaction(request).await?;
 
         Ok(Ok(()))
@@ -439,9 +451,8 @@ impl TaskMgr {
         task_name: &str,
     ) -> Result<Result<Vec<String>, TaskError>, TaskApiError> {
         let task_before_ident = TaskDependentIdent::new_before(&self.tenant, task_name);
-        let task_state_key = TaskStateIdent::new(&self.tenant, task_name);
-        let succeeded_value = TaskState { is_succeeded: true };
-        let not_succeeded_value = TaskState {
+        let succeeded_value = TaskStateValue { is_succeeded: true };
+        let not_succeeded_value = TaskStateValue {
             is_succeeded: false,
         };
 
@@ -460,13 +471,19 @@ impl TaskMgr {
             let Some(target_after_dependent) = target_after_dependent else {
                 continue;
             };
-            let mut request = TxnRequest::new(vec![], vec![])
-                .with_else(vec![txn_put_pb(&task_state_key, &succeeded_value)?]);
-
-            for target_after_task in target_after_dependent.0.iter() {
-                let task_ident = TaskStateIdent::new(&self.tenant, target_after_task);
+            let target_after = &target_after_ident.name().source;
+            let this_task_to_target_state =
+                TaskStateIdent::new(&self.tenant, task_name, target_after);
+            let mut request = TxnRequest::new(vec![], vec![]).with_else(vec![txn_put_pb(
+                &this_task_to_target_state,
+                &succeeded_value,
+            )?]);
+
+            for before_target_after in target_after_dependent.0.iter() {
+                let task_ident =
+                    TaskStateIdent::new(&self.tenant, before_target_after, target_after);
                 // Only care about the predecessors of this task's successor tasks, excluding this task itself.
-                if target_after_task != task_name {
+                if before_target_after != task_name {
                     request.condition.push(TxnCondition::eq_value(
                         task_ident.to_string_key(),
                         succeeded_value.to_pb()?.encode_to_vec(),
@@ -485,15 +502,6 @@ impl TaskMgr {
         Ok(Ok(ready_tasks))
     }
 
-    #[async_backtrace::framed]
-    #[fastrace::trace]
-    pub async fn clean_task_state(&self, task_name: &str) -> Result<(), TaskApiError> {
-        let key = TaskStateIdent::new(&self.tenant, task_name);
-        let req = UpsertPB::delete(key).with(MatchSeq::GE(1));
-        let _ = self.kv_api.upsert_pb(&req).await?;
-        Ok(())
-    }
-
     async fn create_task_inner(
         &self,
         task: Task,